band-aider-pub/band_aider/judgement/unionizing_graph_organizer.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

from band_aider.context.class_meta import ClassMeta
from band_aider.judgement.code_dependency_graph import CodeDependencyGraph
from band_aider.judgement.code_dependency_graph_maintainer import (
    CodeDependencyGraphMaintainer,
    GraphSignature,
)


GraphResolution = Literal["both", "latest", "earliest"]


@dataclass
class UnionizingGraphOrganizer(CodeDependencyGraphMaintainer):
    """UnionizingGraphOrganizer presents multiple maintained graphs as one combined view."""

    inners: tuple[CodeDependencyGraphMaintainer, ...]
    resolution: GraphResolution = "both"
    _graph: CodeDependencyGraph = field(default_factory=CodeDependencyGraph)
    signature: GraphSignature | None = None
    _is_frozen: bool = field(default=False, init=False, repr=False)

    @property
    def graph(self) -> CodeDependencyGraph:
        """Return the current combined graph view."""

        if self._is_frozen:
            return self._graph
        self._graph = self._combine_graphs()
        return self._graph

    @property
    def is_frozen(self) -> bool:
        """Return whether graph mutation is disabled."""

        return self._is_frozen

    def freeze(self) -> None:
        """Freeze the combined graph and every inner maintainer."""

        if self._is_frozen:
            return
        for maintainer in self.inners:
            maintainer.freeze()
        self._graph = self._combine_graphs()
        self._graph.freeze()
        self._is_frozen = True

    def _assert_mutable(self) -> None:
        """Raise when combined graph maintenance mutation is attempted after freeze."""

        if self.is_frozen:
            raise RuntimeError(
                "UnionizingGraphOrganizer is frozen and cannot be mutated."
            )

    def _combine_lookups(
        self,
        attr: str,
    ) -> dict[str, ClassMeta]:
        """Return one merged metadata lookup using the configured conflict policy."""

        merged: dict[str, ClassMeta] = {}
        maintainers = (
            tuple(reversed(self.inners))
            if self.resolution == "earliest"
            else self.inners
        )
        for maintainer in maintainers:
            merged.update(getattr(maintainer.graph, attr))
        return merged

    def _combine_adjacency(
        self,
        attr: str,
    ) -> dict[str, set[str]]:
        """Return one adjacency map unioned across all inner graphs."""

        merged: dict[str, set[str]] = {}
        for maintainer in self.inners:
            for location, edges in getattr(maintainer.graph, attr).items():
                merged.setdefault(location, set()).update(edges)
        return merged

    def _combine_graphs(self) -> CodeDependencyGraph:
        """Return one merged graph from every inner maintainer."""

        type_definitions = {}
        parent_relationships = {}
        declaration_sites = {}
        binding_sites = {}
        return_sites = {}
        indirect_references = {}
        maintainers = (
            tuple(reversed(self.inners))
            if self.resolution == "earliest"
            else self.inners
        )
        for maintainer in maintainers:
            graph = maintainer.graph
            type_definitions.update(graph.type_definitions)
            parent_relationships.update(graph.parent_relationships_by_id)
            declaration_sites.update(graph.declaration_sites_by_id)
            binding_sites.update(graph.binding_sites_by_id)
            return_sites.update(graph.return_sites_by_id)
            indirect_references.update(graph.indirect_references_by_id)
        return CodeDependencyGraph(
            source_by_location=self._combine_lookups("source_by_location"),
            test_by_location=self._combine_lookups("test_by_location"),
            source_dependencies=self._combine_adjacency("source_dependencies"),
            source_dependents=self._combine_adjacency("source_dependents"),
            type_definitions=type_definitions,
            parent_relationships_by_id=parent_relationships,
            declaration_sites_by_id=declaration_sites,
            binding_sites_by_id=binding_sites,
            return_sites_by_id=return_sites,
            indirect_references_by_id=indirect_references,
            tests_for_source=self._combine_adjacency("tests_for_source"),
            sources_for_test=self._combine_adjacency("sources_for_test"),
        )

    def upsert(
        self,
        class_meta: ClassMeta,
        source_files: list[ClassMeta],
        test_files: list[ClassMeta],
        repo_root: Path,
        src_root: Path,
        test_root: Path,
    ) -> CodeDependencyGraph:
        """Upsert one ClassMeta slice in every mutable inner graph."""

        self._assert_mutable()
        for maintainer in self.inners:
            maintainer.upsert(
                class_meta,
                source_files,
                test_files,
                repo_root,
                src_root,
                test_root,
            )
        self._graph = self._combine_graphs()
        return self._graph

    def build(
        self,
        source_files: list[ClassMeta],
        test_files: list[ClassMeta],
        repo_root: Path,
        src_root: Path,
        test_root: Path,
    ) -> CodeDependencyGraph:
        """Rebuild every mutable inner maintainer and return the combined view."""

        self._assert_mutable()
        for maintainer in self.inners:
            maintainer.build(
                source_files,
                test_files,
                repo_root,
                src_root,
                test_root,
            )
        self.signature = (
            tuple(sorted(source.location for source in source_files)),
            tuple(sorted(test.location for test in test_files)),
        )
        return self.graph

    def delete(
        self,
        class_meta: ClassMeta,
        source_files: list[ClassMeta],
        test_files: list[ClassMeta],
        repo_root: Path,
        src_root: Path,
        test_root: Path,
    ) -> CodeDependencyGraph:
        """Delete one ClassMeta slice in every mutable inner graph."""

        self._assert_mutable()
        for maintainer in self.inners:
            maintainer.delete(
                class_meta,
                source_files,
                test_files,
                repo_root,
                src_root,
                test_root,
            )
        self._graph = self._combine_graphs()
        return self._graph