mirror of
https://github.com/hanzoai/aci.git
synced 2026-11-02 19:07:28 +00:00
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
from typing import Dict, Optional, Set
|
|
|
|
from networkx import DiGraph
|
|
from pydantic import (
|
|
ConfigDict,
|
|
)
|
|
|
|
from .codeblocks import (
|
|
BlockSpan,
|
|
CodeBlock,
|
|
CodeBlockType,
|
|
SpanType,
|
|
)
|
|
|
|
|
|
class Module(CodeBlock):
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
|
|
file_path: str
|
|
content: str
|
|
spans_by_id: Dict[str, BlockSpan] = {}
|
|
language: str | None = None
|
|
parent: CodeBlock | None = None
|
|
|
|
_graph: DiGraph = None # TODO: Move to central CodeGraph
|
|
|
|
def __init__(self, **data):
|
|
data.setdefault('type', CodeBlockType.MODULE)
|
|
super().__init__(**data)
|
|
|
|
def find_span_by_id(self, span_id: str) -> Optional[BlockSpan]:
|
|
return self.spans_by_id.get(span_id)
|
|
|
|
def sum_tokens(self, span_ids: Optional[Set[str]] = None):
|
|
tokens = self.tokens
|
|
if span_ids:
|
|
for span_id in span_ids:
|
|
span = self.spans_by_id.get(span_id)
|
|
if span:
|
|
tokens += span.tokens
|
|
return tokens
|
|
|
|
tokens += sum([child.sum_tokens() for child in self.children])
|
|
return tokens
|
|
|
|
def find_related_span_ids(self, span_id: str) -> Set[str]:
|
|
related_span_ids = set()
|
|
|
|
blocks = self.find_blocks_by_span_id(span_id)
|
|
for block in blocks:
|
|
# Find successors (outgoing relationships)
|
|
successors = list(self._graph.successors(block.path_string()))
|
|
for succ in successors:
|
|
node_data = self._graph.nodes[succ]
|
|
if 'block' in node_data:
|
|
span = node_data['block'].belongs_to_span
|
|
related_span_ids.add(span.span_id)
|
|
|
|
# Find predecessors (incoming relationships)
|
|
predecessors = list(self._graph.predecessors(block.path_string()))
|
|
for pred in predecessors:
|
|
node_data = self._graph.nodes[pred]
|
|
if 'block' in node_data:
|
|
span = node_data['block'].belongs_to_span
|
|
related_span_ids.add(span.span_id)
|
|
|
|
# Always add parent class initation span
|
|
if (
|
|
block.parent
|
|
and block.parent.type == CodeBlockType.CLASS
|
|
and block.belongs_to_span is not None
|
|
):
|
|
related_span_ids.add(block.belongs_to_span.span_id)
|
|
for class_child in block.parent.children:
|
|
if (
|
|
class_child.belongs_to_span
|
|
and class_child.belongs_to_span.span_type == SpanType.INITATION
|
|
):
|
|
related_span_ids.add(class_child.belongs_to_span.span_id)
|
|
|
|
# Always add module initation span
|
|
for span in self.spans_by_id.values():
|
|
if (
|
|
span.block_type == CodeBlockType.MODULE
|
|
and span.span_type == SpanType.INITATION
|
|
):
|
|
related_span_ids.add(span.span_id)
|
|
|
|
return related_span_ids
|