Class ModelAnalyzer

utils/generate_model_tests.py:79–193

class ModelAnalyzer(ast.NodeVisitor):
    def __init__(self):
        self.model_classes = []
        self.current_class = None
        self.imports = set()

    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            self.imports.add(alias.name.split(".")[-1])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom):
        for alias in node.names:
            self.imports.add(alias.name)
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef):
        base_names = []
        for base in node.bases:
            if isinstance(base, ast.Name):
                base_names.append(base.id)
            elif isinstance(base, ast.Attribute):
                base_names.append(base.attr)

        if "ModelMixin" in base_names:
            class_info = {
                "name": node.name,
                "bases": base_names,
                "attributes": {},
                "has_forward": False,
                "init_params": [],
            }

            for item in node.body:
                if isinstance(item, ast.Assign):
                    for target in item.targets:
                        if isinstance(target, ast.Name):
                            attr_name = target.id
                            if attr_name.startswith("_"):
                                class_info["attributes"][attr_name] = self._get_value(item.value)

                elif isinstance(item, ast.FunctionDef):
                    if item.name == "forward":
                        class_info["has_forward"] = True
                        class_info["forward_params"] = self._extract_func_params(item)
                    elif item.name == "__init__":
                        class_info["init_params"] = self._extract_func_params(item)

            self.model_classes.append(class_info)

        self.generic_visit(node)

    def _extract_func_params(self, func_node: ast.FunctionDef) -> list[dict]:
        params = []
        args = func_node.args

        num_defaults = len(args.defaults)
        num_args = len(args.args)