Fixed assignment of objects and added tests

lcompilers · Jul 15, 2024 · 285b965 · 285b965
1 parent c01688f
commit 285b965
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/integration_tests/class_02.py b/integration_tests/class_02.py
@@ -10,16 +10,15 @@ def attack(self:"Character", other:"Character") -> str:
         other.health -= self.attack_power
         return self.name+" attacks "+ other.name+" for "+str(self.attack_power)+" damage."
 
-
     def is_alive(self:"Character")->bool:
         if self.is_immortal:
             return True
         else:
             return self.health > 0
+
 def main():
     hero : Character = Character("Hero", 10, 20)
     monster : Character = Character("Monster", 50, 15)
-
     print(hero.attack(monster))
     print(monster.health)  
     assert monster.health == 30
@@ -36,5 +35,10 @@ def main():
     hero.is_immortal = False
     print(hero.is_alive())
     assert hero.is_alive() == False
+    print("Restarting")
+    hero = Character("Hero", 10, 20)
+    print(hero.is_alive())
+    assert hero.is_alive() == True
 
 main()
+
diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -5549,6 +5549,26 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
             }
             tmp_vec.push_back(ASR::make_Assignment_t(al, x.base.base.loc, target, tmp_value,
                                     overloaded));
+            if ( target->type == ASR::exprType::Var &&
+                    tmp_value->type == ASR::exprType::StructConstructor ) {
+                Vec<ASR::call_arg_t> new_args; new_args.reserve(al, 1);
+                ASR::call_arg_t self_arg;
+                self_arg.loc = x.base.base.loc;
+                ASR::symbol_t* st = ASR::down_cast<ASR::Var_t>(target)->m_v;
+                self_arg.m_value = target;
+                new_args.push_back(al,self_arg);
+                AST::Call_t* call = AST::down_cast<AST::Call_t>(x.m_value);
+                if ( call->n_keywords>0 ) {
+                    throw SemanticError("Kwargs not implemented yet", x.base.base.loc);
+                }
+                visit_expr_list(call->m_args, call->n_args, new_args);
+                ASR::symbol_t* der = ASR::down_cast<ASR::StructType_t>(
+                    ASR::down_cast<ASR::Variable_t>(st)->m_type)->m_derived_type;
+                std::string call_name = "__init__";
+                ASR::symbol_t* call_sym = get_struct_member(der, call_name, x.base.base.loc);
+                tmp_vec.push_back(make_call_helper(al, call_sym,
+                    current_scope, new_args, call_name, x.base.base.loc));
+            }
         }
         // to make sure that we add only those statements in tmp_vec
         tmp = nullptr;