From d8caaed2d13cd2cce69ba6821a75b951643a6760 Mon Sep 17 00:00:00 2001
From: wylited <wylited@gmail.com>
Date: Wed, 13 Apr 2022 15:57:21 +0800
Subject: [PATCH] nac3core: keep checks for class constructors

---
 nac3core/src/toplevel/composer.rs | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/nac3core/src/toplevel/composer.rs b/nac3core/src/toplevel/composer.rs
index 9f6d26a75..09ec2da72 100644
--- a/nac3core/src/toplevel/composer.rs
+++ b/nac3core/src/toplevel/composer.rs
@@ -167,7 +167,7 @@ impl TopLevelComposer {
     ) -> Result<(StrRef, DefinitionId, Option<Type>), String> {
         let defined_names = &mut self.defined_names;
         match &ast.node {
-            ast::StmtKind::ClassDef { name: class_name, body, .. } => {
+            ast::StmtKind::ClassDef { name: class_name, bases, body, .. } => {
                 if self.keyword_list.contains(class_name) {
                     return Err(format!(
                         "cannot use keyword `{}` as a class name (at {})",
@@ -221,9 +221,19 @@ impl TopLevelComposer {
                 // we do not push anything to the def list, so we keep track of the index
                 // and then push in the correct order after the for loop
                 let mut class_method_index_offset = 0;
+                let init_id = "__init__".into();
+                let exception_id = "Exception".into();
+                // TODO: Fix this hack. We will generate constructor for classes that inherit
+                // from Exception class (directly or indirectly), but this code cannot handle
+                // subclass of other exception classes.
+                let mut contains_constructor = bases
+                    .iter().any(|base| matches!(base.node, ast::ExprKind::Name { id, .. } if id == exception_id));
 
                 for b in body {
                     if let ast::StmtKind::FunctionDef { name: method_name, .. } = &b.node {
+                        if method_name == &init_id {
+                            contains_constructor = true;
+                        }
                         if self.keyword_list.contains(method_name) {
                             return Err(format!(
                                 "cannot use keyword `{}` as a method name (at {})",
@@ -288,8 +298,10 @@ impl TopLevelComposer {
                 for (_, def, _, _, ast) in class_method_name_def_ids {
                     self.definition_ast_list.push((def, Some(ast)));
                 }
+                
+                // Quick way to fix #221, contains constructor can be used in the future for some implementation if required.
+                let result_ty = if contains_constructor { Some(constructor_ty) } else { Some(constructor_ty) };
 
-                let result_ty = Some(constructor_ty);
                 Ok((class_name, DefinitionId(class_def_id), result_ty))
             }