fix: use ast to create funix at the runtime

backend/funix/decorator/__init__.py

@@ -1,6 +1,7 @@
 """
 Funix decorator. The central logic of Funix.
 """
+import ast
 import dataclasses
 import inspect
 import sys
@@ -14,20 +15,22 @@
     Parameter,
     Signature,
     getsource,
+    isclass,
     isgeneratorfunction,
     ismethod,
     signature,
-    isclass,
 )
 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
 from urllib.request import urlopen
 from uuid import uuid4
 
+import dill
 from flask import Response, request, session
 from requests import post
 from requests.structures import CaseInsensitiveDict
@@ -57,6 +60,7 @@
     get_type_dict,
     get_type_widget_prop,
 )
+from funix.decorator.runtime import RuntimeClassVisitor
 from funix.hint import (
     ArgumentConfigType,
     ConditionalVisibleType,
@@ -76,8 +80,8 @@
 )
 from funix.session import (
     get_global_variable,
-    set_global_variable,
     set_default_global_variable,
+    set_global_variable,
 )
 from funix.theme import get_dict_theme, parse_theme
 from funix.util.module import funix_menu_to_safe_function_name
@@ -1808,6 +1812,7 @@ def wrapped_function(**wrapped_function_kwargs):
                                 "error_body": format_exc(),
                             }
 
+                    @wraps(function)
                     def output_to_web_function(**wrapped_function_kwargs):
                         try:
                             fake_stdout = StdoutToWebsocket(ws)
@@ -2013,83 +2018,51 @@ def output_to_web_function(**wrapped_function_kwargs):
     return decorator
 
 
-def funix_class(cls):
-    if isclass(cls):
-        if not hasattr(cls, "__init__"):
-            raise Exception("Class must have `__init__` method!")
-
-        set_default_global_variable("__FUNIX_" + cls.__name__, None)
-
-        init_class = signature(cls)
-
-        if len(init_class.parameters) == 0:
-
-            def create_class():
-                set_global_variable("__FUNIX_" + cls.__name__, cls())
-
-        else:
-
-            def create_class(*replace_args, **replace_kwargs):
-                set_global_variable(
-                    "__FUNIX_" + cls.__name__, cls(*replace_args, **replace_kwargs)
-                )
-
-            create_class.__signature__ = signature(cls)
-
-        setattr(create_class, "__name__", cls.__name__)
-
-        funix()(create_class)
-
-        def get_init_function():
-            inited_class = get_global_variable("__FUNIX_" + cls.__name__)
-
-            if inited_class is None:
-                raise Exception("Class must be inited first!")
+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:
-                return inited_class
-
-        for class_function in dir(cls):
-            if not class_function.startswith("_"):
-                function = getattr(cls, class_function)
-                if callable(function):
-                    is_static = isinstance(
-                        inspect.getattr_static(cls, function.__name__), staticmethod
-                    )
-
-                    sign = signature(function)
-
-                    is_empty = (
-                        len(sign.parameters) == 0
-                        if is_static
-                        else len(sign.parameters) == 1
-                    )
+                inside_class = False
 
-                    if is_empty:
+    return dedent("".join(class_code)).strip()
 
-                        def new_change_function():
-                            if is_static:
-                                return function()
-                            api = get_init_function()
-                            return function(api)
 
-                    else:
-
-                        def new_change_function(*function_args, **function_kwargs):
-                            if is_static:
-                                return function(*function_args, **function_kwargs)
-                            api = get_init_function()
-                            return function(api, *function_args, **function_kwargs)
-
-                    if is_static:
-                        new_change_function.__signature__ = sign
-                    else:
-                        new_sign = sign.replace(
-                            parameters=tuple(sign.parameters.values())[1:]
-                        )
-                        new_change_function.__signature__ = new_sign
+def funix_class(cls):
+    if inspect.isclass(cls):
+        if not hasattr(cls, "__init__"):
+            raise Exception("Class must have __init__ method!")
 
-                    setattr(new_change_function, "__name__", function.__name__)
-                    funix()(new_change_function)
+        f = RuntimeClassVisitor(cls.__name__, funix, cls)
+        f.visit(
+            ast.parse(
+                get_class_source_code(inspect.getsourcefile(cls.__init__), cls.__name__)
+            )
+        )
     else:
         for class_function in dir(cls):
             if not class_function.startswith("_"):

backend/funix/decorator/runtime.py

@@ -0,0 +1,161 @@
+import ast
+from _ast import (
+    Assign,
+    Attribute,
+    Call,
+    ClassDef,
+    Constant,
+    Expr,
+    FunctionDef,
+    Load,
+    Module,
+    Name,
+    Return,
+    Store,
+    keyword,
+)
+from typing import Any
+
+from funix.session import get_global_variable, set_global_variable
+
+
+def get_init_function(cls_name: str):
+    inited_class = get_global_variable("__FUNIX_" + cls_name)
+
+    if inited_class is None:
+        raise Exception("Class must be inited first!")
+    else:
+        return inited_class
+
+
+def set_init_function(cls_name: str, cls):
+    set_global_variable("__FUNIX_" + cls_name, cls)
+
+
+class RuntimeClassVisitor(ast.NodeVisitor):
+    def __init__(self, cls_name: str, funix: Any, cls: Any):
+        self.funix = funix
+        self._cls_name = cls_name
+        self._cls = cls
+
+    def visit_ClassDef(self, node: ClassDef) -> Any:
+        for cls_function in node.body:
+            if isinstance(cls_function, FunctionDef):
+                self.visit_FunctionDef(cls_function)
+
+    def visit_FunctionDef(self, node: FunctionDef) -> Any:
+        args = node.args
+
+        is_static_method = False
+
+        for decorator in node.decorator_list:
+            if hasattr(decorator, "id") and decorator.id == "staticmethod":
+                is_static_method = True
+                break
+
+            # TODO: handle from funix import funix_class_params as xxx
+            #       Or xxx = funix_class_params
+            funix_decorator = Call(
+                func=Name(id="funix", ctx=Load()), args=[], keywords=[]
+            )
+            if hasattr(decorator, "func") and decorator.func.id == "funix_class_params":
+                if (
+                    decorator.keywords[0].arg == "disable"
+                    and decorator.keywords[0].value.value
+                ):
+                    return
+                else:
+                    funix_decorator.keywords = decorator.keywords
+                    funix_decorator.args = decorator.args
+
+        if not is_static_method:
+            # Yes .args
+            args.args.pop(0)
+
+        if node.name.startswith("_") and node.name != "__init__":
+            return
+
+        new_module = Module(
+            body=[
+                FunctionDef(
+                    name=node.name,
+                    args=args,
+                    decorator_list=[
+                        Call(func=Name(id="funix", ctx=Load()), args=[], keywords=[])
+                    ],
+                    returns=node.returns,
+                    lineno=0,
+                )
+            ],
+            type_ignores=[],
+        )
+
+        if node.name == "__init__":
+            new_module.body[0].decorator_list[0].keywords = [
+                keyword(arg="title", value=Constant(value=self._cls_name))
+            ]
+            new_module.body[0].name = "new_" + self._cls_name
+            new_module.body[0].body = [
+                Expr(
+                    value=Call(
+                        func=Name(id="set_init_function", ctx=Load()),
+                        args=[
+                            Constant(value=self._cls_name),
+                            Call(
+                                func=Name(id=self._cls_name, ctx=Load()),
+                                args=[
+                                    Name(id=args.arg, ctx=Load()) for args in args.args
+                                ],
+                                keywords=[],
+                            ),
+                        ],
+                        keywords=[],
+                    )
+                )
+            ]
+        else:
+            if is_static_method:
+                new_module.body[0].body = [
+                    Return(
+                        value=Call(
+                            func=Attribute(
+                                value=Name(id=f"{self._cls_name}", ctx=Load()),
+                                attr=node.name,
+                                ctx=Load(),
+                            ),
+                            args=[Name(id=arg.arg, ctx=Load()) for arg in args.args],
+                            keywords=[],
+                        )
+                    )
+                ]
+            else:
+                new_module.body[0].body = [
+                    Assign(
+                        targets=[Name(id="api", ctx=Store())],
+                        value=Call(
+                            func=Name(id="get_init_function", ctx=Load()),
+                            args=[Constant(value=self._cls_name)],
+                            keywords=[],
+                        ),
+                        lineno=0,
+                    ),
+                    Return(
+                        value=Call(
+                            func=Attribute(
+                                value=Name(id="api", ctx=Load()),
+                                attr=node.name,
+                                ctx=Load(),
+                            ),
+                            args=[Name(id=args.arg, ctx=Load()) for args in args.args],
+                            keywords=[],
+                        )
+                    ),
+                ]
+        new_globals = globals()
+        new_globals["funix"] = self.funix
+        new_globals[self._cls_name] = self._cls
+        exec(
+            ast.unparse(new_module),
+            globals(),
+            locals(),
+        )

examples/class.py

@@ -1,4 +1,4 @@
-from funix import funix_class, funix_class_params
+from funix import funix_class
 
 
 class A: