From 0391158e6683b8bdb4f371851418e8e252e651b9 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 15 May 2026 00:48:08 +0800 Subject: [PATCH] up --- reme2/component/file_graph/__init__.py | 4 +- .../component/file_graph/local_file_graph.py | 128 ++++++++++++++++++ 2 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 reme2/component/file_graph/local_file_graph.py diff --git a/reme2/component/file_graph/__init__.py b/reme2/component/file_graph/__init__.py index a1fa3e8c..bbb1bbf1 100644 --- a/reme2/component/file_graph/__init__.py +++ b/reme2/component/file_graph/__init__.py @@ -1,11 +1,13 @@ """File graph module.""" from .base_file_graph import BaseFileGraph -from .nx_file_graph import NxFileGraph +from .local_file_graph import LocalFileGraph from .neo4j_file_graph import Neo4jFileGraph +from .nx_file_graph import NxFileGraph __all__ = [ "BaseFileGraph", + "LocalFileGraph", "NxFileGraph", "Neo4jFileGraph", ] diff --git a/reme2/component/file_graph/local_file_graph.py b/reme2/component/file_graph/local_file_graph.py new file mode 100644 index 00000000..c2eb74b9 --- /dev/null +++ b/reme2/component/file_graph/local_file_graph.py @@ -0,0 +1,128 @@ +"""Pure-Python file-graph backend (no external deps).""" + +from pathlib import Path + +from .base_file_graph import BaseFileGraph +from ..component_registry import R +from ...schema import FileLink, FileNode + + +@R.register("local") +class LocalFileGraph(BaseFileGraph): + """Dict-backed file graph; trusts ``FileLink.path`` for adjacency.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._nodes: dict[str, FileNode] = {} + self._inverse: dict[str, set[str]] = {} + # Virtual edges: src→target where target isn't in ``_nodes`` yet. + self._pending: dict[str, set[str]] = {} + self._graph_file: Path = self.graph_path / f"{self.graph_name}.jsonl" + + # -- Lifecycle --------------------------------------------------------- + + async def _start(self) -> None: + await super()._start() + self._load() + await self.rebuild_links() + edges = sum(len(s) for s in self._inverse.values()) + pending = sum(len(s) for s in self._pending.values()) + self.logger.info( + f"LocalFileGraph '{self.graph_name}' ready: " + f"{len(self._nodes)} nodes, {edges} edges, {pending} pending", + ) + + async def _close(self) -> None: + self._dump() + await super()._close() + + def _load(self) -> None: + if not self._graph_file.exists(): + return + with open(self._graph_file, "r", encoding="utf-8") as f: + self._nodes.update( + (n.path, n) + for n in (FileNode.model_validate_json(line) for line in f if line.strip()) + ) + + def _dump(self) -> None: + tmp = self._graph_file.with_suffix(".tmp") + with open(tmp, "w", encoding="utf-8") as f: + f.writelines(f"{n.model_dump_json()}\n" for n in self._nodes.values()) + tmp.replace(self._graph_file) + + # -- Edge bookkeeping -------------------------------------------------- + + def _add_edge(self, src: str, target: str) -> None: + bucket = self._inverse if target in self._nodes else self._pending + bucket.setdefault(target, set()).add(src) + + def _remove_edge(self, src: str, target: str) -> None: + for bucket in (self._inverse, self._pending): + srcs = bucket.get(target) + if srcs is None or src not in srcs: + continue + srcs.discard(src) + if not srcs: + del bucket[target] + + # -- Node CRUD --------------------------------------------------------- + + async def upsert_nodes(self, nodes: list[FileNode]) -> None: + for node in nodes: + path = node.path + old = self._nodes.get(path) + if old is not None: + for link in old.links: + if link.path: + self._remove_edge(path, link.path) + self._nodes[path] = node + for link in node.links: + if link.path: + self._add_edge(path, link.path) + # Promote virtual edges aimed at this newly-arrived target. + promoted = self._pending.pop(path, None) + if promoted: + self._inverse.setdefault(path, set()).update(promoted) + + async def delete_nodes(self, paths: list[str]) -> None: + for path in paths: + node = self._nodes.pop(path, None) + if node is None: + continue + for link in node.links: + if link.path: + self._remove_edge(path, link.path) + # Demote inbound edges to virtual; sources still link here. + demoted = self._inverse.pop(path, None) + if demoted: + self._pending.setdefault(path, set()).update(demoted) + + async def get_nodes(self, paths: list[str]) -> list[FileNode]: + return [self._nodes[p] for p in paths if p in self._nodes] + + async def rebuild_links(self) -> None: + self._inverse.clear() + self._pending.clear() + for src, node in self._nodes.items(): + for link in node.links: + if link.path: + self._add_edge(src, link.path) + + # -- Link access ------------------------------------------------------- + + async def get_outlinks(self, path: str) -> list[FileLink]: + node = self._nodes.get(path) + if node is None: + return [] + return [link for link in node.links if link.path and link.path in self._nodes] + + async def get_inlinks(self, path: str) -> list[FileLink]: + if path not in self._nodes: + return [] + return [ + link + for src in self._inverse.get(path, ()) + for link in self._nodes[src].links + if link.path == path + ] \ No newline at end of file