app/services/parser_service.py

import ast
from pathlib import Path

from app.models.code import (
    ImportInfo,
    FunctionInfo,
    ClassInfo,
    PythonFileInfo,
    CallInfo,
)


class FunctionCallVisitor(ast.NodeVisitor):

    def __init__(self):
        self.calls = []

    def visit_Call(self, node: ast.Call):

        name = self._get_call_name(node.func)

        if name is not None:

            self.calls.append(
                CallInfo(
                    name=name,
                    line=node.lineno,
                )
            )

        self.generic_visit(node)

    def _get_call_name(self, node):

        if isinstance(node, ast.Name):
            return node.id

        if isinstance(node, ast.Attribute):

            parts = []

            while isinstance(node, ast.Attribute):
                parts.append(node.attr)
                node = node.value

            if isinstance(node, ast.Name):
                parts.append(node.id)

            return ".".join(reversed(parts))

        return None


class ParserService:
    def _module_name(self, file_path: Path):

        parts = list(file_path.parts)

        if "app" in parts:

            index = parts.index("app")

            parts = parts[index:]

        parts[-1] = parts[-1].replace(".py", "")

        return ".".join(parts)

    def parse_file(self, file_path: Path) -> PythonFileInfo:

        module_name = self._module_name(file_path)

        source = file_path.read_text(
            encoding="utf-8",
            errors="ignore",
        )

        lines = source.splitlines()

        tree = ast.parse(source)

        imports = []
        classes = []
        functions = []

        for node in tree.body:

            if isinstance(node, ast.Import):

                for alias in node.names:

                    imports.append(
                        ImportInfo(
                            module=alias.name,
                            name=alias.name,
                            alias=alias.asname,
                        )
                    )

            elif isinstance(node, ast.ImportFrom):

                for alias in node.names:

                    imports.append(
                        ImportInfo(
                            module=node.module,
                            name=alias.name,
                            alias=alias.asname,
                        )
                    )

            elif isinstance(node, ast.FunctionDef):

                functions.append(
                    self._parse_function(
                        node,
                        module_name,
                        lines=lines,
                    )
                )

            elif isinstance(node, ast.ClassDef):

                methods = []

                for child in node.body:

                    if isinstance(child, ast.FunctionDef):

                        methods.append(
                            self._parse_function(
                                node=child,
                                module_name=module_name,
                                lines=lines,
                                class_name=node.name,
                            )
                        )
                class_source = self._get_source_code(
                    node,
                    lines,
                )

                classes.append(

                    ClassInfo(

                        name=node.name,

                        qualified_name=f"{module_name}.{node.name}",

                        line=node.lineno,

                        end_line=node.end_lineno,

                        source_code=class_source,

                        methods=methods,
                    )
                )

        return PythonFileInfo(
            path=str(file_path),
            imports=imports,
            classes=classes,
            functions=functions,
        )
    def _get_source_code(
        self,
        node,
        lines,
    ) -> str:

        return "\n".join(
            lines[node.lineno - 1 : node.end_lineno]
        )

    def _parse_function(
        self,
        node,
        module_name,
        lines,
        class_name=None,
    ):

        visitor = FunctionCallVisitor()

        visitor.visit(node)

        if class_name:

            qualified = f"{module_name}.{class_name}.{node.name}"

        else:

            qualified = f"{module_name}.{node.name}"

        print("module_name:", repr(module_name))
        print("class_name:", repr(class_name))
        print("node.name:", repr(node.name))
        print("qualified:", repr(qualified))        
        source_code = self._get_source_code(
            node,
            lines,
        )

        return FunctionInfo(

            name=node.name,

            qualified_name=qualified,

            line=node.lineno,

            end_line=node.end_lineno,

            docstring=ast.get_docstring(node),

            source_code=source_code,

            calls=visitor.calls,
        )