295 lines
10 KiB
Python
295 lines
10 KiB
Python
# 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 band_aider.context.class_meta import ClassMeta
|
|
from band_aider.judgement.code_dependency_graph_maintainer import (
|
|
CodeDependencyGraphMaintainer,
|
|
)
|
|
from band_aider.judgement.code_dependency_graph import DependencySelector
|
|
from band_aider.judgement.dependency_discovery_strategy import (
|
|
DependencyDirection,
|
|
DependencyDiscoveryStrategy,
|
|
DependencyTraversalCache,
|
|
_dependency_traversal_query,
|
|
_stable_quicksort_sources_by_upstream,
|
|
)
|
|
from band_aider.judgement.python_ast_dependency_graph_maintainer import (
|
|
PythonAstDependencyGraphMaintainer,
|
|
)
|
|
|
|
@dataclass
|
|
class PythonAstDependencyStrategy(DependencyDiscoveryStrategy):
|
|
"""Answer dependency queries through one maintained Python AST graph.
|
|
|
|
Example: ``strategy.source_dependents_of("src/join.py")`` returns direct
|
|
dependents from the graph this strategy owns.
|
|
"""
|
|
|
|
maintainer: CodeDependencyGraphMaintainer = field(
|
|
default_factory=PythonAstDependencyGraphMaintainer
|
|
)
|
|
|
|
@property
|
|
def graph(self):
|
|
"""Return the maintained dependency graph."""
|
|
|
|
return self.maintainer.graph
|
|
|
|
@property
|
|
def is_frozen(self) -> bool:
|
|
"""Return whether the underlying maintainer has been frozen."""
|
|
|
|
return self.maintainer.is_frozen
|
|
|
|
def freeze(self) -> None:
|
|
"""Freeze the underlying maintainer."""
|
|
|
|
self.maintainer.freeze()
|
|
|
|
def source_dependencies_of(
|
|
self,
|
|
source_location: str,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
) -> set[str]:
|
|
"""Return one-hop prerequisites from this strategy's maintained graph."""
|
|
|
|
return self.maintainer.source_dependencies_of(
|
|
source_location,
|
|
selector=selector,
|
|
)
|
|
|
|
def source_dependents_of(
|
|
self,
|
|
source_location: str,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
) -> set[str]:
|
|
"""Return one-hop dependents from this strategy's maintained graph."""
|
|
|
|
return self.maintainer.source_dependents_of(
|
|
source_location,
|
|
selector=selector,
|
|
)
|
|
|
|
def has_path(
|
|
self,
|
|
start: str,
|
|
target: str,
|
|
direction: DependencyDirection = DependencyDirection.UPSTREAM,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
excluded_locations: frozenset[str] = frozenset(),
|
|
) -> bool:
|
|
"""Return whether a source path reaches its target before crossing excluded relays."""
|
|
|
|
# Reject a path whose starting relay has already been declared dead.
|
|
if start in excluded_locations:
|
|
return False
|
|
# Admit the identity path without traversing the graph.
|
|
if start == target:
|
|
return True
|
|
|
|
step = (
|
|
self.maintainer.source_dependencies_of
|
|
if direction is DependencyDirection.UPSTREAM
|
|
else self.maintainer.source_dependents_of
|
|
)
|
|
queue = [start]
|
|
seen = {start}
|
|
# Traverse only until the target is found, preserving predicate semantics.
|
|
while queue:
|
|
location = queue.pop(0)
|
|
# Inspect one-hop neighbors while protecting cycles and dead relays.
|
|
for neighbor in sorted(step(location, selector=selector)):
|
|
# Stop immediately when the requested connectivity is witnessed.
|
|
if neighbor == target:
|
|
return True
|
|
# Skip dead relays and previously inspected locations.
|
|
if neighbor in excluded_locations or neighbor in seen:
|
|
continue
|
|
seen.add(neighbor)
|
|
queue.append(neighbor)
|
|
return False
|
|
|
|
def direct_tests_for(
|
|
self,
|
|
source_files: list[ClassMeta],
|
|
all_test_files: list[ClassMeta] | None = None,
|
|
src_root: Path | None = None,
|
|
test_root: Path | None = None,
|
|
) -> list[ClassMeta]:
|
|
"""Return direct tests from the dependency graph."""
|
|
|
|
return self.maintainer.direct_tests_for(source_files)
|
|
|
|
def _walk_source_frontier(
|
|
self,
|
|
source_files: list[ClassMeta],
|
|
step,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
) -> list[ClassMeta]:
|
|
"""Return source metadata reached by repeatedly applying one shallow source step."""
|
|
|
|
queue = list(step(source_files, selector=selector))
|
|
traversed: dict[str, ClassMeta] = {
|
|
source.location: source
|
|
for source in queue
|
|
}
|
|
while queue:
|
|
source = queue.pop(0)
|
|
for next_source in step([source], selector=selector):
|
|
if next_source.location in traversed:
|
|
continue
|
|
traversed[next_source.location] = next_source
|
|
queue.append(next_source)
|
|
return [
|
|
traversed[location]
|
|
for location in sorted(traversed)
|
|
]
|
|
|
|
def downstream_sources_for(
|
|
self,
|
|
source_files: list[ClassMeta],
|
|
all_source_files: list[ClassMeta] | None = None,
|
|
src_root: Path | None = None,
|
|
recursive: bool = False,
|
|
dependency_cache: DependencyTraversalCache | None = None,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
) -> list[ClassMeta]:
|
|
"""Return downstream sources from the dependency graph."""
|
|
|
|
query = _dependency_traversal_query(
|
|
source_files,
|
|
recursive=recursive,
|
|
selector=selector,
|
|
)
|
|
if dependency_cache is not None:
|
|
cached = dependency_cache.downstream.get(query)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
if recursive:
|
|
result = self._walk_source_frontier(
|
|
source_files,
|
|
self.maintainer.downstream_sources_for,
|
|
selector,
|
|
)
|
|
else:
|
|
result = self.maintainer.downstream_sources_for(source_files, selector=selector)
|
|
|
|
if dependency_cache is not None:
|
|
dependency_cache.downstream[query] = tuple(result)
|
|
return list(result)
|
|
|
|
def upstream_sources_for(
|
|
self,
|
|
files: list[ClassMeta],
|
|
all_source_files: list[ClassMeta] | None = None,
|
|
src_root: Path | None = None,
|
|
test_root: Path | None = None,
|
|
recursive: bool = False,
|
|
dependency_cache: DependencyTraversalCache | None = None,
|
|
selector: DependencySelector = DependencySelector.BOTH,
|
|
) -> list[ClassMeta]:
|
|
"""Return upstream source context for source or test files from the dependency graph."""
|
|
|
|
query = _dependency_traversal_query(
|
|
files,
|
|
recursive=recursive,
|
|
selector=selector,
|
|
)
|
|
if dependency_cache is not None:
|
|
cached = dependency_cache.upstream.get(query)
|
|
if cached is not None:
|
|
return list(cached)
|
|
|
|
if recursive:
|
|
result = self._walk_source_frontier(
|
|
files,
|
|
self.maintainer.upstream_sources_for,
|
|
selector,
|
|
)
|
|
else:
|
|
result = self.maintainer.upstream_sources_for(files, selector=selector)
|
|
if dependency_cache is not None:
|
|
dependency_cache.upstream[query] = tuple(result)
|
|
return list(result)
|
|
|
|
def rerank_sources_by_upstream(
|
|
self,
|
|
source_files: list[ClassMeta],
|
|
all_source_files: list[ClassMeta],
|
|
src_root: Path,
|
|
test_root: Path,
|
|
dependency_cache: DependencyTraversalCache,
|
|
) -> list[ClassMeta]:
|
|
"""Stable-quicksort roots by cached recursive upstream membership."""
|
|
|
|
unique_roots: list[ClassMeta] = []
|
|
seen: set[str] = set()
|
|
for source_file in source_files:
|
|
if source_file.location in seen:
|
|
continue
|
|
seen.add(source_file.location)
|
|
unique_roots.append(source_file)
|
|
|
|
upstream_locations_by_location: dict[str, frozenset[str]] = {}
|
|
for root in unique_roots:
|
|
upstream_locations_by_location[root.location] = frozenset(
|
|
source.location
|
|
for source in self.upstream_sources_for(
|
|
[root],
|
|
all_source_files,
|
|
src_root,
|
|
test_root,
|
|
recursive=True,
|
|
dependency_cache=dependency_cache,
|
|
)
|
|
)
|
|
return _stable_quicksort_sources_by_upstream(
|
|
unique_roots,
|
|
upstream_locations_by_location,
|
|
)
|
|
|
|
def joint_beam_sources_for(
|
|
self,
|
|
current_frontier: list[ClassMeta],
|
|
deferred_blames: list[ClassMeta],
|
|
all_source_files: list[ClassMeta] | None = None,
|
|
src_root: Path | None = None,
|
|
) -> list[ClassMeta]:
|
|
"""Return graph intersection candidates for deferred-guided expansion."""
|
|
|
|
current_locations = {
|
|
source.location
|
|
for source in current_frontier
|
|
}
|
|
deferred_locations = {
|
|
source.location
|
|
for source in deferred_blames
|
|
}
|
|
downstream_sources = self._walk_source_frontier(
|
|
current_frontier,
|
|
self.maintainer.downstream_sources_for,
|
|
)
|
|
upstream_sources = self._walk_source_frontier(
|
|
deferred_blames,
|
|
self.maintainer.upstream_sources_for,
|
|
)
|
|
downstream = current_locations | {source.location for source in downstream_sources}
|
|
upstream_to_deferred = {source.location for source in upstream_sources}
|
|
upstream_to_deferred.update(deferred_locations)
|
|
candidates = downstream & upstream_to_deferred
|
|
sources_by_location = {
|
|
source.location: source
|
|
for source in [*downstream_sources, *upstream_sources, *deferred_blames]
|
|
}
|
|
return [
|
|
sources_by_location[location]
|
|
for location in sorted(candidates)
|
|
if location in sources_by_location
|
|
]
|