fix: remove get_class_source_code, use full ast analyse

backend/funix/decorator/__init__.py

@@ -11,17 +11,10 @@
 from enum import Enum, auto
 from functools import wraps
 from importlib import import_module
-from inspect import (
-    Parameter,
-    Signature,
-    getsource,
-    isgeneratorfunction,
-    signature,
-)
+from inspect import Parameter, Signature, getsource, isgeneratorfunction, signature
 from io import StringIO
 from json import dumps, loads
 from secrets import token_hex
-from textwrap import dedent
 from traceback import format_exc
 from types import ModuleType
 from typing import Any, Optional
@@ -2011,40 +2004,6 @@ def output_to_web_function(**wrapped_function_kwargs):
     return decorator
 
 
-def get_class_source_code(file_path: str, class_name: str) -> str:
-    with open(file_path, "r") as f:
-        lines = f.readlines()
-
-    inside_class = False
-    is_now_next = False
-    class_code = []
-
-    hope_indent = 0
-
-    for line in lines:
-        if line.strip().startswith("class " + class_name) and line.strip().endswith(
-            ":"
-        ):
-            inside_class = True
-            is_now_next = True
-            class_code.append(line)
-        elif inside_class:
-            if line.strip() == "":
-                class_code.append(line)
-                continue
-            if line.strip() and not line.startswith(" "):
-                inside_class = False
-            if is_now_next:
-                hope_indent = len(line) - len(line.lstrip())
-                is_now_next = False
-            if line.startswith(" " * hope_indent):
-                class_code.append(line)
-            else:
-                inside_class = False
-
-    return dedent("".join(class_code)).strip()
-
-
 def funix_class():
     return __funix_class
 
@@ -2055,11 +2014,11 @@ def __funix_class(cls):
             raise Exception("Class must have __init__ method!")
 
         f = RuntimeClassVisitor(cls.__name__, funix, cls)
-        f.visit(
-            ast.parse(
-                get_class_source_code(inspect.getsourcefile(cls.__init__), cls.__name__)
-            )
-        )
+
+        with open(inspect.getsourcefile(cls.__init__), "r") as file_:
+            class_source_code = file_.read()
+
+        f.visit(ast.parse(class_source_code))
         return cls
     else:
         for class_function in dir(cls):

backend/funix/decorator/runtime.py

@@ -1,5 +1,4 @@
 import ast
-import inspect
 from _ast import (
     Assign,
     Attribute,
@@ -33,32 +32,22 @@ def set_init_function(cls_name: str, cls):
     set_global_variable("__FUNIX_" + cls_name, cls)
 
 
-def get_imports(path) -> list:
-    with open(path) as fh:
-        root = ast.parse(fh.read(), path)
-
-    imports = []
-
-    for node in ast.iter_child_nodes(root):
-        if isinstance(node, ast.Import):
-            imports.append(node)
-        elif isinstance(node, ast.ImportFrom):
-            imports.append(node)
-        else:
-            continue
-
-    return imports
-
-
 class RuntimeClassVisitor(ast.NodeVisitor):
-    def __init__(self, cls_name: str, funix: Any, cls: Any):
-        self.funix = funix
+    def __init__(self, cls_name: str, funix_: Any, cls: Any):
+        self.funix = funix_
         self._cls_name = cls_name
         self._cls = cls
-        source_path = inspect.getsourcefile(self._cls.__init__)
-        self._imports = get_imports(source_path)
+        self._imports = []
+
+    def visit_Import(self, node):
+        self._imports.append(node)
+
+    def visit_ImportFrom(self, node):
+        self._imports.append(node)
 
     def visit_ClassDef(self, node: ClassDef) -> Any:
+        if node.name != self._cls_name:
+            return
         for cls_function in node.body:
             if isinstance(cls_function, FunctionDef):
                 self.visit_FunctionDef(cls_function)