app/services/graph_service.py

from app.repositories.neo4j_repository import Neo4jRepository


class GraphService:

    def __init__(self):

        self.neo = Neo4jRepository()

    def build(
        self,
        repository_id: str,
        resolved,
    ):

        for file in resolved:

            self.neo.merge_file(
                repository_id=repository_id,
                path=file["path"],
            )

            self._create_imports(
                repository_id,
                file,
            )

            self._create_functions(
                repository_id,
                file,
            )

            self._create_classes(
                repository_id,
                file,
            )

            self._create_declares(
                repository_id,
                file,
            )

            self._create_has_methods(
                repository_id,
                file,
            )

            self._create_edges(
                repository_id,
                file,
            )

        self.neo.create_vector_indexes()

    def _create_functions(
        self,
        repository_id: str,
        file,
    ):

        for function in file["functions"]:

            self.neo.merge_function(
                repository_id=repository_id,
                qualified_name=function["qualified_name"],
                name=function["name"],
                source_code=function["source_code"],
            )

    def _create_classes(
        self,
        repository_id: str,
        file,
    ):

        for cls in file["classes"]:

            self.neo.merge_class(
                repository_id=repository_id,
                qualified_name=cls["qualified_name"],
                name=cls["name"],
                source_code=cls["source_code"],
            )

            for method in cls["methods"]:

                self.neo.merge_function(
                    repository_id=repository_id,
                    qualified_name=method["qualified_name"],
                    name=method["name"],
                    source_code=method["source_code"],
                )

    def _create_declares(
        self,
        repository_id: str,
        file,
    ):

        for function in file["functions"]:

            self.neo.create_declares(
                repository_id=repository_id,
                file_path=file["path"],
                qualified_name=function["qualified_name"],
            )

        for cls in file["classes"]:

            self.neo.create_declares(
                repository_id=repository_id,
                file_path=file["path"],
                qualified_name=cls["qualified_name"],
            )

    def _create_has_methods(
        self,
        repository_id: str,
        file,
    ):

        for cls in file["classes"]:

            for method in cls["methods"]:

                self.neo.create_has_method(
                    repository_id=repository_id,
                    class_name=cls["qualified_name"],
                    method_name=method["qualified_name"],
                )

    def _create_imports(
        self,
        repository_id: str,
        file,
    ):

        for imp in file["imports"]:

            module = imp["module"]

            if not module:
                continue

            self.neo.merge_module(
                repository_id=repository_id,
                module=module,
            )

            self.neo.create_import(
                repository_id=repository_id,
                file_path=file["path"],
                module=module,
                symbol=imp["name"],
                alias=imp["alias"],
            )

    def _create_edges(
        self,
        repository_id: str,
        file,
    ):

        self._create_function_edges(
            repository_id,
            file["functions"],
        )

        for cls in file["classes"]:

            self._create_function_edges(
                repository_id,
                cls["methods"],
            )

    def _create_function_edges(
        self,
        repository_id: str,
        functions,
    ):

        for function in functions:

            source = function["qualified_name"]

            for call in function["calls"]:

                target = call["qualified_name"]

                if target is None:
                    continue

                if target.startswith("app."):

                    self.neo.merge_call(
                        repository_id=repository_id,
                        source=source,
                        target=target,
                    )

                else:

                    self.neo.merge_external_function(
                        repository_id=repository_id,
                        qualified_name=target,
                    )

                    self.neo.create_belongs_to(
                        repository_id=repository_id,
                        qualified_name=target,
                    )

                    self.neo.merge_call(
                        repository_id=repository_id,
                        source=source,
                        target=target,
                    )