                if int(os.path.getmtime(f)) != int(mt.get(f, 0)):
                    return False
                if self._file_hash(f) != str(hh.get(f, "")):
                    return False
            # Load
            self.indexed_files = files
            for s in data.get("symbols", []):
                sym = Symbol(**s)
                self._add_symbol(sym)
            for mod, mi in data.get("modules", {}).items():
                self.modules[mod] = ModuleInfo(**mi)
            self.calls = [tuple(x) for x in data.get("calls", [])]
            self.module_to_tests = data.get("module_to_tests", {})
            self.module_imports = data.get("module_imports", {})
            self._cached_mtimes = {k: int(v) for k, v in (mt or {}).items()}
            self._cached_hashes = {k: str(v) for k, v in (hh or {}).items()}
            return True
        except Exception:
            return False

    def _load_cache_relaxed(self, cache_path: str) -> bool:
        try:
            if not os.path.exists(cache_path):
                return False
            data = json.loads(open(cache_path, "r", encoding="utf-8").read())
            if str(data.get("version", "")) != "3":
                return False
            self.indexed_files = data.get("indexed_files", [])
            for s in data.get("symbols", []):
                sym = Symbol(**s)
                self._add_symbol(sym)
            for mod, mi in data.get("modules", {}).items():
                self.modules[mod] = ModuleInfo(**mi)
            self.calls = [tuple(x) for x in data.get("calls", [])]
            self.module_to_tests = data.get("module_to_tests", {})
            self.module_imports = data.get("module_imports", {})
            self._cached_mtimes = {
                k: int(v) for k, v in (data.get("mtimes", {}) or {}).items()
            }
            self._cached_hashes = {
                k: str(v) for k, v in (data.get("hashes", {}) or {}).items()
            }
            return True
        except Exception:
            return False

    def _detect_changed_files(
        self, old_mt: Dict[str, int], old_hh: Dict[str, str]
    ) -> Tuple[List[str], List[str]]:
        curr_files: List[str] = []
        for dirpath, dirnames, filenames in os.walk(self.root):
            dir_rel = os.path.relpath(dirpath, self.root)
            dirnames[:] = [d for d in dirnames if not self._is_ignored(os.path.join(dir_rel, d))]
            if self._is_ignored(dir_rel):
                continue
            for fn in filenames:
                if fn.endswith(".py"):
                    fp = os.path.join(dirpath, fn)
                    if self._is_ignored(os.path.relpath(fp, self.root)):
                        continue
                    curr_files.append(fp)
        curr = set(curr_files)
        prev = set(self.indexed_files or [])
        removed = list(prev - curr)
        added = list(curr - prev)


                return False
            self.indexed_files = data.get("indexed_files", [])
            for s in data.get("symbols", []):
                sym = Symbol(**s)
                self._add_symbol(sym)
            for mod, mi in data.get("modules", {}).items():
                self.modules[mod] = ModuleInfo(**mi)
            self.calls = [tuple(x) for x in data.get("calls", [])]
            self.module_to_tests = data.get("module_to_tests", {})
            self.module_imports = data.get("module_imports", {})
            self._cached_mtimes = {
                k: int(v) for k, v in (data.get("mtimes", {}) or {}).items()
            }
            self._cached_hashes = {
                k: str(v) for k, v in (data.get("hashes", {}) or {}).items()
            }
            return True
        except Exception:
            return False

    def _detect_changed_files(
        self, old_mt: Dict[str, int], old_hh: Dict[str, str]
    ) -> Tuple[List[str], List[str]]:
        curr_files: List[str] = []
        for dirpath, dirnames, filenames in os.walk(self.root):
            dir_rel = os.path.relpath(dirpath, self.root)
            dirnames[:] = [d for d in dirnames if not self._is_ignored(os.path.join(dir_rel, d))]
            if self._is_ignored(dir_rel):
                continue
            for fn in filenames:
                if fn.endswith(".py"):
                    fp = os.path.join(dirpath, fn)
                    if self._is_ignored(os.path.relpath(fp, self.root)):
                        continue
                    curr_files.append(fp)
        curr = set(curr_files)
        prev = set(self.indexed_files or [])
        removed = list(prev - curr)
        added = list(curr - prev)
        changed: List[str] = list(added)
        for f in curr & prev:
            try:
                mt = int(os.path.getmtime(f))
                hh = self._file_hash(f)
            except Exception:
                changed.append(f)
                continue
            if old_mt.get(f) != mt or old_hh.get(f) != hh:
                changed.append(f)
        self.indexed_files = sorted(list(curr))
        return sorted(set(changed)), sorted(removed)

    def _incremental_reindex(
        self, changed_files: List[str], removed_files: List[str]
    ) -> None:
        # purge removed modules
        for f in removed_files:
            m = self._module_name_for_path(f)
            if m in self.modules:
                to_remove = [
                    fqn for fqn, s in list(self.symbols_by_fqn.items()) if s.module == m
                ]
                for fqn in to_remove:
                    s = self.symbols_by_fqn.pop(fqn, None)
                    if s:
                        self.symbols_by_name[s.name] = [
                            x for x in self.symbols_by_name.get(s.name, []) if x != fqn
                        ]
                self.calls = [
                    (a, b) for (a, b) in self.calls if not a.startswith(m + ".")
                ]


                    if self._is_ignored(os.path.relpath(fp, self.root)):
                        continue
                    curr_files.append(fp)
        curr = set(curr_files)
        prev = set(self.indexed_files or [])
        removed = list(prev - curr)
        added = list(curr - prev)
        changed: List[str] = list(added)
        for f in curr & prev:
            try:
                mt = int(os.path.getmtime(f))
                hh = self._file_hash(f)
            except Exception:
                changed.append(f)
                continue
            if old_mt.get(f) != mt or old_hh.get(f) != hh:
                changed.append(f)
        self.indexed_files = sorted(list(curr))
        return sorted(set(changed)), sorted(removed)

    def _incremental_reindex(
        self, changed_files: List[str], removed_files: List[str]
    ) -> None:
        # purge removed modules
        for f in removed_files:
            m = self._module_name_for_path(f)
            if m in self.modules:
                to_remove = [
                    fqn for fqn, s in list(self.symbols_by_fqn.items()) if s.module == m
                ]
                for fqn in to_remove:
                    s = self.symbols_by_fqn.pop(fqn, None)
                    if s:
                        self.symbols_by_name[s.name] = [
                            x for x in self.symbols_by_name.get(s.name, []) if x != fqn
                        ]
                self.calls = [
                    (a, b) for (a, b) in self.calls if not a.startswith(m + ".")
                ]
                self.modules.pop(m, None)
                self.module_imports.pop(m, None)
                self.module_star_imports.pop(m, None)
        mods = {self._module_name_for_path(f) for f in changed_files}
        # include dependents via reverse import graph
        rev = self._reverse_imports()
        queue = list(mods)
        seen = set(mods)
        while queue:
            m = queue.pop(0)
            for dep in rev.get(m, []):
                if dep not in seen:
                    seen.add(dep)
                    queue.append(dep)
        for m in seen:
            self._reindex_module(m)

    def _reverse_imports(self) -> Dict[str, List[str]]:
        rev: Dict[str, List[str]] = {}
        for m, deps in self.module_imports.items():
            for d in deps:
                rev.setdefault(d, []).append(m)
        return rev

    def _reindex_module(self, module: str) -> None:
        mi = self.modules.get(module)
        if not mi:
            return
        # remove existing symbols and calls for this module
        to_remove = [
            fqn for fqn, s in self.symbols_by_fqn.items() if s.module == module
        ]
        for fqn in to_remove:
            s = self.symbols_by_fqn.pop(fqn, None)
            if s:
                lst = self.symbols_by_name.get(s.name, [])


                self.calls = [
                    (a, b) for (a, b) in self.calls if not a.startswith(m + ".")
                ]
                self.modules.pop(m, None)
                self.module_imports.pop(m, None)
                self.module_star_imports.pop(m, None)
        mods = {self._module_name_for_path(f) for f in changed_files}
        # include dependents via reverse import graph
        rev = self._reverse_imports()
        queue = list(mods)
        seen = set(mods)
        while queue:
            m = queue.pop(0)
            for dep in rev.get(m, []):
                if dep not in seen:
                    seen.add(dep)
                    queue.append(dep)
        for m in seen:
            self._reindex_module(m)

    def _reverse_imports(self) -> Dict[str, List[str]]:
        rev: Dict[str, List[str]] = {}
        for m, deps in self.module_imports.items():
            for d in deps:
                rev.setdefault(d, []).append(m)
        return rev

    def _reindex_module(self, module: str) -> None:
        mi = self.modules.get(module)
        if not mi:
            return
        # remove existing symbols and calls for this module
        to_remove = [
            fqn for fqn, s in self.symbols_by_fqn.items() if s.module == module
        ]
        for fqn in to_remove:
            s = self.symbols_by_fqn.pop(fqn, None)
            if s:
                lst = self.symbols_by_name.get(s.name, [])
                self.symbols_by_name[s.name] = [x for x in lst if x != fqn]
        self.calls = [(a, b) for (a, b) in self.calls if not a.startswith(module + ".")]
        # reset import maps for this module
        self.modules[module].imports = {}
        self.module_imports[module] = []
        self.module_star_imports[module] = []
        # re-parse


        # include dependents via reverse import graph
        rev = self._reverse_imports()
        queue = list(mods)
        seen = set(mods)
        while queue:
            m = queue.pop(0)
            for dep in rev.get(m, []):
                if dep not in seen:
                    seen.add(dep)
                    queue.append(dep)
        for m in seen:
            self._reindex_module(m)

    def _reverse_imports(self) -> Dict[str, List[str]]:
        rev: Dict[str, List[str]] = {}
        for m, deps in self.module_imports.items():
            for d in deps:
                rev.setdefault(d, []).append(m)
        return rev

    def _reindex_module(self, module: str) -> None:
        mi = self.modules.get(module)
        if not mi:
            return
        # remove existing symbols and calls for this module
        to_remove = [
            fqn for fqn, s in self.symbols_by_fqn.items() if s.module == module
        ]
        for fqn in to_remove:
            s = self.symbols_by_fqn.pop(fqn, None)
            if s:
                lst = self.symbols_by_name.get(s.name, [])
                self.symbols_by_name[s.name] = [x for x in lst if x != fqn]
        self.calls = [(a, b) for (a, b) in self.calls if not a.startswith(module + ".")]
        # reset import maps for this module
        self.modules[module].imports = {}
        self.module_imports[module] = []
        self.module_star_imports[module] = []
        # re-parse
        try:
            src = open(mi.file, "r", encoding="utf-8").read()
            tree = ast.parse(src)
        except Exception:
            return
        visitor = _ModuleVisitor(module, mi.file)
        visitor.visit(tree)
        self.modules[module].imports.update(visitor.imports)
        self.module_imports[module] = sorted(visitor.import_modules)
        self.module_star_imports[module] = list(getattr(visitor, "star_imports", []))
        self.modules[module].exports = list(getattr(visitor, "exports", []))
        for sym in visitor.symbols:
            self._add_symbol(sym)
        for caller, callee_key in visitor.calls:
            callee_fqn = self._resolve_callee(module, callee_key, visitor)
            self.calls.append((caller, callee_fqn or callee_key))

    def _save_cache(self, cache_path: str) -> None:
        try:
            mt = {f: int(os.path.getmtime(f)) for f in self.indexed_files}
            hh = {f: self._file_hash(f) for f in self.indexed_files}
            data = {
                "version": "3",
                "indexed_files": self.indexed_files,
                "mtimes": mt,
                "hashes": hh,
                "symbols": [self._sym_to_dict(s) for s in self.symbols_by_fqn.values()],
                "modules": {k: self._mi_to_dict(v) for k, v in self.modules.items()},
                "calls": self.calls,
                "module_to_tests": self.module_to_tests,
                "module_imports": self.module_imports,
            }
            open(cache_path, "w", encoding="utf-8").write(json.dumps(data))
        except Exception:
            pass


        self.module_imports[module] = []
        self.module_star_imports[module] = []
        # re-parse
        try:
            src = open(mi.file, "r", encoding="utf-8").read()
            tree = ast.parse(src)
        except Exception:
            return
        visitor = _ModuleVisitor(module, mi.file)
        visitor.visit(tree)
        self.modules[module].imports.update(visitor.imports)
        self.module_imports[module] = sorted(visitor.import_modules)
        self.module_star_imports[module] = list(getattr(visitor, "star_imports", []))
        self.modules[module].exports = list(getattr(visitor, "exports", []))
        for sym in visitor.symbols:
            self._add_symbol(sym)
        for caller, callee_key in visitor.calls:
            callee_fqn = self._resolve_callee(module, callee_key, visitor)
            self.calls.append((caller, callee_fqn or callee_key))

    def _save_cache(self, cache_path: str) -> None:
        try:
            mt = {f: int(os.path.getmtime(f)) for f in self.indexed_files}
            hh = {f: self._file_hash(f) for f in self.indexed_files}
            data = {
                "version": "3",
                "indexed_files": self.indexed_files,
                "mtimes": mt,
                "hashes": hh,
                "symbols": [self._sym_to_dict(s) for s in self.symbols_by_fqn.values()],
                "modules": {k: self._mi_to_dict(v) for k, v in self.modules.items()},
                "calls": self.calls,
                "module_to_tests": self.module_to_tests,
                "module_imports": self.module_imports,
            }
            open(cache_path, "w", encoding="utf-8").write(json.dumps(data))
        except Exception:
            pass

    def _sym_to_dict(self, s: Symbol) -> Dict[str, Any]:
        return {
            "fqn": s.fqn,
            "name": s.name,
            "qualname": s.qualname,
            "kind": s.kind,
            "module": s.module,
            "file": s.file,
            "line": s.line,
            "end_line": s.end_line,
            "doc": s.doc,
            "signature": s.signature,
            "returns": s.returns,
        }

    def _mi_to_dict(self, mi: ModuleInfo) -> Dict[str, Any]:
        return {
            "module": mi.module,
            "file": mi.file,


    def _save_cache(self, cache_path: str) -> None:
        try:
            mt = {f: int(os.path.getmtime(f)) for f in self.indexed_files}
            hh = {f: self._file_hash(f) for f in self.indexed_files}
            data = {
                "version": "3",
                "indexed_files": self.indexed_files,
                "mtimes": mt,
                "hashes": hh,
                "symbols": [self._sym_to_dict(s) for s in self.symbols_by_fqn.values()],
                "modules": {k: self._mi_to_dict(v) for k, v in self.modules.items()},
                "calls": self.calls,
                "module_to_tests": self.module_to_tests,
                "module_imports": self.module_imports,
            }
            open(cache_path, "w", encoding="utf-8").write(json.dumps(data))
        except Exception:
            pass

    def _sym_to_dict(self, s: Symbol) -> Dict[str, Any]:
        return {
            "fqn": s.fqn,
            "name": s.name,
            "qualname": s.qualname,
            "kind": s.kind,
            "module": s.module,
            "file": s.file,
            "line": s.line,
            "end_line": s.end_line,
            "doc": s.doc,
            "signature": s.signature,
            "returns": s.returns,
        }

    def _mi_to_dict(self, mi: ModuleInfo) -> Dict[str, Any]:
        return {
            "module": mi.module,
            "file": mi.file,
            "is_test": mi.is_test,
            "imports": mi.imports,
            "defs": mi.defs,
            "exports": mi.exports,
        }

    def _is_builtin_name(self, name: str) -> bool:
        try:
            import builtins as _bi  # type: ignore

            if hasattr(_bi, name):
                return True
        except Exception:
            pass
        return name in {


            }
            open(cache_path, "w", encoding="utf-8").write(json.dumps(data))
        except Exception:
            pass

    def _sym_to_dict(self, s: Symbol) -> Dict[str, Any]:
        return {
            "fqn": s.fqn,
            "name": s.name,
            "qualname": s.qualname,
            "kind": s.kind,
            "module": s.module,
            "file": s.file,
            "line": s.line,
            "end_line": s.end_line,
            "doc": s.doc,
            "signature": s.signature,
            "returns": s.returns,
        }

    def _mi_to_dict(self, mi: ModuleInfo) -> Dict[str, Any]:
        return {
            "module": mi.module,
            "file": mi.file,
            "is_test": mi.is_test,
            "imports": mi.imports,
            "defs": mi.defs,
            "exports": mi.exports,
        }

    def _is_builtin_name(self, name: str) -> bool:
        try:
            import builtins as _bi  # type: ignore

            if hasattr(_bi, name):
                return True
        except Exception:
            pass
        return name in {
            "super",
            "property",
            "globals",
            "locals",
            "__import__",
            "print",
            "len",
            "range",
            "dict",
            "list",


            "kind": s.kind,
            "module": s.module,
            "file": s.file,
            "line": s.line,
            "end_line": s.end_line,
            "doc": s.doc,
            "signature": s.signature,
            "returns": s.returns,
        }

    def _mi_to_dict(self, mi: ModuleInfo) -> Dict[str, Any]:
        return {
            "module": mi.module,
            "file": mi.file,
            "is_test": mi.is_test,
            "imports": mi.imports,
            "defs": mi.defs,
            "exports": mi.exports,
        }

    def _is_builtin_name(self, name: str) -> bool:
        try:
            import builtins as _bi  # type: ignore

            if hasattr(_bi, name):
                return True
        except Exception:
            pass
        return name in {
            "super",
            "property",
            "globals",
            "locals",
            "__import__",
            "print",
            "len",
            "range",
            "dict",
            "list",
            "set",
            "tuple",
            "int",
            "float",
            "bool",
            "max",
            "min",
            "sum",
            "open",
            "enumerate",
            "zip",
            "map",
            "filter",
            "round",
            "any",
            "all",
            "sorted",
            "hasattr",
            "getattr",
            "setattr",
            "isinstance",
            "issubclass",
        }

    def _file_hash(self, path: str) -> str:
        try:
            import hashlib

            with open(path, "rb") as rf:
                return hashlib.sha1(rf.read()).hexdigest()
        except Exception:
            return ""

    # --- Coverage --- #

    def attach_coverage_from_xml(self, xml_path: str) -> None:
        try:
            import xml.etree.ElementTree as ET

            tree = ET.parse(xml_path)
            root = tree.getroot()
            files_hits: Dict[str, set[int]] = {}
            # coverage.py XML: <class filename="path"> ... <lines><line number="N" hits="H"/></lines>


            "bool",
            "max",
            "min",
            "sum",
            "open",
            "enumerate",
            "zip",
            "map",
            "filter",
            "round",
            "any",
            "all",
            "sorted",
            "hasattr",
            "getattr",
            "setattr",
            "isinstance",
            "issubclass",
        }

    def _file_hash(self, path: str) -> str:
        try:
            import hashlib

            with open(path, "rb") as rf:
                return hashlib.sha1(rf.read()).hexdigest()
        except Exception:
            return ""

    # --- Coverage --- #

    def attach_coverage_from_xml(self, xml_path: str) -> None:
        try:
            import xml.etree.ElementTree as ET

            tree = ET.parse(xml_path)
            root = tree.getroot()
            files_hits: Dict[str, set[int]] = {}
            # coverage.py XML: <class filename="path"> ... <lines><line number="N" hits="H"/></lines>
            for cls in root.findall(".//class"):
                fn = cls.attrib.get("filename", "")
                if not fn:
                    continue
                # Normalize to absolute path
                f_abs = (
                    fn
                    if os.path.isabs(fn)
                    else os.path.abspath(os.path.join(self.root, fn))


            "all",
            "sorted",
            "hasattr",
            "getattr",
            "setattr",
            "isinstance",
            "issubclass",
        }

    def _file_hash(self, path: str) -> str:
        try:
            import hashlib

            with open(path, "rb") as rf:
                return hashlib.sha1(rf.read()).hexdigest()
        except Exception:
            return ""

    # --- Coverage --- #

    def attach_coverage_from_xml(self, xml_path: str) -> None:
        try:
            import xml.etree.ElementTree as ET

            tree = ET.parse(xml_path)
            root = tree.getroot()
            files_hits: Dict[str, set[int]] = {}
            # coverage.py XML: <class filename="path"> ... <lines><line number="N" hits="H"/></lines>
            for cls in root.findall(".//class"):
                fn = cls.attrib.get("filename", "")
                if not fn:
                    continue
                # Normalize to absolute path
                f_abs = (
                    fn
                    if os.path.isabs(fn)
                    else os.path.abspath(os.path.join(self.root, fn))
                )
                hits = files_hits.setdefault(f_abs, set())
                for ln in cls.findall(".//line"):
                    try:
                        num = int(ln.attrib.get("number", "0"))
                        h = int(ln.attrib.get("hits", "0"))
                        if h > 0:
                            hits.add(num)
                    except Exception:
                        continue
            # Some coverage.xml variants place <file> nodes
            if not files_hits:
                for fnode in root.findall(".//file"):
                    fn = fnode.attrib.get("filename", "")
                    if not fn:
                        continue
                    f_abs = (
                        fn
                        if os.path.isabs(fn)
                        else os.path.abspath(os.path.join(self.root, fn))
                    )
                    hits = files_hits.setdefault(f_abs, set())
                    for ln in fnode.findall(".//line"):
                        try:
                            num = int(ln.attrib.get("number", "0"))
                            h = int(ln.attrib.get("hits", "0"))
                            if h > 0:
                                hits.add(num)
                        except Exception:
                            continue
            self.coverage_files = files_hits
            # Compute per-symbol coverage
            sym_cov: Dict[str, float] = {}
            for fqn, sym in self.symbols_by_fqn.items():
                covered = files_hits.get(sym.file, set())
                a = int(sym.line)
                b = int(sym.end_line) if int(sym.end_line) >= a else a
                span = list(range(a, b + 1))
                if not span:
                    sym_cov[fqn] = 0.0
                    continue
                hits = sum(1 for x in span if x in covered)
                sym_cov[fqn] = hits / float(len(span))
            self.symbol_coverage = sym_cov
        except Exception:
            # Leave coverage empty on error
            self.coverage_files = {}
            self.symbol_coverage = {}

    def coverage_of(self, fqn: str) -> Optional[float]:
        return self.symbol_coverage.get(fqn)

def _cli() -> None:
    import argparse
    import json

    p = argparse.ArgumentParser()
    p.add_argument("root", nargs="?", default="./repo")
    p.add_argument("--ignore", action="append", default=None, help="Relative paths to ignore (repeatable)")
    p.add_argument("--owners-of", dest="owners_of", default=None)
    p.add_argument("--search", dest="search", default=None)
    p.add_argument("--defs-in", dest="defs_in", default=None)
    p.add_argument("--calls-of", dest="calls_of", default=None)
    p.add_argument("--who-calls", dest="who_calls", default=None)
    p.add_argument("--dump", dest="dump", action="store_true")
    p.add_argument("--coverage-xml", dest="coverage_xml", default=None)
    p.add_argument("--coverage-of", dest="coverage_of", default=None)


                            continue
            self.coverage_files = files_hits
            # Compute per-symbol coverage
            sym_cov: Dict[str, float] = {}
            for fqn, sym in self.symbols_by_fqn.items():
                covered = files_hits.get(sym.file, set())
                a = int(sym.line)
                b = int(sym.end_line) if int(sym.end_line) >= a else a
                span = list(range(a, b + 1))
                if not span:
                    sym_cov[fqn] = 0.0
                    continue
                hits = sum(1 for x in span if x in covered)
                sym_cov[fqn] = hits / float(len(span))
            self.symbol_coverage = sym_cov
        except Exception:
            # Leave coverage empty on error
            self.coverage_files = {}
            self.symbol_coverage = {}

    def coverage_of(self, fqn: str) -> Optional[float]:
        return self.symbol_coverage.get(fqn)

def _cli() -> None:
    import argparse
    import json

    p = argparse.ArgumentParser()
    p.add_argument("root", nargs="?", default="./repo")
    p.add_argument("--ignore", action="append", default=None, help="Relative paths to ignore (repeatable)")
    p.add_argument("--owners-of", dest="owners_of", default=None)
    p.add_argument("--search", dest="search", default=None)
    p.add_argument("--defs-in", dest="defs_in", default=None)
    p.add_argument("--calls-of", dest="calls_of", default=None)
    p.add_argument("--who-calls", dest="who_calls", default=None)
    p.add_argument("--dump", dest="dump", action="store_true")
    p.add_argument("--coverage-xml", dest="coverage_xml", default=None)
    p.add_argument("--coverage-of", dest="coverage_of", default=None)
    p.add_argument("--refs-of", dest="refs_of", default=None)
    p.add_argument("--tests-for", dest="tests_for", default=None)
    p.add_argument("--tests-for-module", dest="tests_for_module", default=None)


            )
        )
        return
    # Dump summary
    print(json.dumps({"files": len(g.indexed_files), "symbols": len(g.symbols_by_fqn)}))

class _ModuleVisitor(ast.NodeVisitor):
    def __init__(self, module: str, path: str) -> None:
        self.module = module
        self.path = path
        self.symbols: List[Symbol] = []
        self.calls: List[Tuple[str, str]] = []  # (caller_fqn, callee_key)
        self.stack: List[str] = []  # qualname stack
        self.class_stack: List[str] = []
        self.imports: Dict[str, str] = {}
        self.import_modules: List[str] = []
        self.star_imports: List[str] = []
        self.exports: List[str] = []

    def _cur_qualname(self) -> str:
        return ".".join(self.stack)

    def _cur_class(self) -> Optional[str]:
        return self.class_stack[-1] if self.class_stack else None

    def _fqn(self, name: str) -> str:
        q = self._cur_qualname()
        return f"{self.module}.{q + ('.' if q else '')}{name}"

    def visit_Import(self, node: ast.Import) -> Any:  # type: ignore[override]
        for alias in node.names:
            asname = alias.asname or alias.name.split(".")[-1]
            self.imports[asname] = alias.name
            self.import_modules.append(alias.name.split(".")[0])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:  # type: ignore[override]
        # Resolve relative imports: from .x import y
        if node.level and node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))


    # Dump summary
    print(json.dumps({"files": len(g.indexed_files), "symbols": len(g.symbols_by_fqn)}))

class _ModuleVisitor(ast.NodeVisitor):
    def __init__(self, module: str, path: str) -> None:
        self.module = module
        self.path = path
        self.symbols: List[Symbol] = []
        self.calls: List[Tuple[str, str]] = []  # (caller_fqn, callee_key)
        self.stack: List[str] = []  # qualname stack
        self.class_stack: List[str] = []
        self.imports: Dict[str, str] = {}
        self.import_modules: List[str] = []
        self.star_imports: List[str] = []
        self.exports: List[str] = []

    def _cur_qualname(self) -> str:
        return ".".join(self.stack)

    def _cur_class(self) -> Optional[str]:
        return self.class_stack[-1] if self.class_stack else None

    def _fqn(self, name: str) -> str:
        q = self._cur_qualname()
        return f"{self.module}.{q + ('.' if q else '')}{name}"

    def visit_Import(self, node: ast.Import) -> Any:  # type: ignore[override]
        for alias in node.names:
            asname = alias.asname or alias.name.split(".")[-1]
            self.imports[asname] = alias.name
            self.import_modules.append(alias.name.split(".")[0])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:  # type: ignore[override]
        # Resolve relative imports: from .x import y
        if node.level and node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            prefix = base[:-up] if up > 0 else base
            mod = ".".join([p for p in prefix if p] + [node.module])
        elif node.level and not node.module:


class _ModuleVisitor(ast.NodeVisitor):
    def __init__(self, module: str, path: str) -> None:
        self.module = module
        self.path = path
        self.symbols: List[Symbol] = []
        self.calls: List[Tuple[str, str]] = []  # (caller_fqn, callee_key)
        self.stack: List[str] = []  # qualname stack
        self.class_stack: List[str] = []
        self.imports: Dict[str, str] = {}
        self.import_modules: List[str] = []
        self.star_imports: List[str] = []
        self.exports: List[str] = []

    def _cur_qualname(self) -> str:
        return ".".join(self.stack)

    def _cur_class(self) -> Optional[str]:
        return self.class_stack[-1] if self.class_stack else None

    def _fqn(self, name: str) -> str:
        q = self._cur_qualname()
        return f"{self.module}.{q + ('.' if q else '')}{name}"

    def visit_Import(self, node: ast.Import) -> Any:  # type: ignore[override]
        for alias in node.names:
            asname = alias.asname or alias.name.split(".")[-1]
            self.imports[asname] = alias.name
            self.import_modules.append(alias.name.split(".")[0])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:  # type: ignore[override]
        # Resolve relative imports: from .x import y
        if node.level and node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            prefix = base[:-up] if up > 0 else base
            mod = ".".join([p for p in prefix if p] + [node.module])
        elif node.level and not node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            mod = ".".join(base[:-up])
        else:


        self.path = path
        self.symbols: List[Symbol] = []
        self.calls: List[Tuple[str, str]] = []  # (caller_fqn, callee_key)
        self.stack: List[str] = []  # qualname stack
        self.class_stack: List[str] = []
        self.imports: Dict[str, str] = {}
        self.import_modules: List[str] = []
        self.star_imports: List[str] = []
        self.exports: List[str] = []

    def _cur_qualname(self) -> str:
        return ".".join(self.stack)

    def _cur_class(self) -> Optional[str]:
        return self.class_stack[-1] if self.class_stack else None

    def _fqn(self, name: str) -> str:
        q = self._cur_qualname()
        return f"{self.module}.{q + ('.' if q else '')}{name}"

    def visit_Import(self, node: ast.Import) -> Any:  # type: ignore[override]
        for alias in node.names:
            asname = alias.asname or alias.name.split(".")[-1]
            self.imports[asname] = alias.name
            self.import_modules.append(alias.name.split(".")[0])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:  # type: ignore[override]
        # Resolve relative imports: from .x import y
        if node.level and node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            prefix = base[:-up] if up > 0 else base
            mod = ".".join([p for p in prefix if p] + [node.module])
        elif node.level and not node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            mod = ".".join(base[:-up])
        else:
            mod = node.module or ""
        for alias in node.names:
            # star import
            if getattr(alias, "name", "") == "*":
                if mod:
                    self.star_imports.append(mod)
                continue


        self.star_imports: List[str] = []
        self.exports: List[str] = []

    def _cur_qualname(self) -> str:
        return ".".join(self.stack)

    def _cur_class(self) -> Optional[str]:
        return self.class_stack[-1] if self.class_stack else None

    def _fqn(self, name: str) -> str:
        q = self._cur_qualname()
        return f"{self.module}.{q + ('.' if q else '')}{name}"

    def visit_Import(self, node: ast.Import) -> Any:  # type: ignore[override]
        for alias in node.names:
            asname = alias.asname or alias.name.split(".")[-1]
            self.imports[asname] = alias.name
            self.import_modules.append(alias.name.split(".")[0])
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:  # type: ignore[override]
        # Resolve relative imports: from .x import y
        if node.level and node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            prefix = base[:-up] if up > 0 else base
            mod = ".".join([p for p in prefix if p] + [node.module])
        elif node.level and not node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            mod = ".".join(base[:-up])
        else:
            mod = node.module or ""
        for alias in node.names:
            # star import
            if getattr(alias, "name", "") == "*":
                if mod:
                    self.star_imports.append(mod)
                continue
            asname = alias.asname or alias.name
            self.imports[asname] = f"{mod}.{alias.name}" if mod else alias.name
        if mod:
            self.import_modules.append(mod.split(".")[0])
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> Any:  # type: ignore[override]
        fqn = self._fqn(node.name)
        try:
            doc_s = ast.get_docstring(node) or None
        except Exception:
            doc_s = None
        sym = Symbol(
            fqn=fqn,
            name=node.name,
            qualname=self._cur_qualname(),
            kind="class",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)


            prefix = base[:-up] if up > 0 else base
            mod = ".".join([p for p in prefix if p] + [node.module])
        elif node.level and not node.module:
            base = self.module.split(".")
            up = max(0, int(node.level))
            mod = ".".join(base[:-up])
        else:
            mod = node.module or ""
        for alias in node.names:
            # star import
            if getattr(alias, "name", "") == "*":
                if mod:
                    self.star_imports.append(mod)
                continue
            asname = alias.asname or alias.name
            self.imports[asname] = f"{mod}.{alias.name}" if mod else alias.name
        if mod:
            self.import_modules.append(mod.split(".")[0])
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> Any:  # type: ignore[override]
        fqn = self._fqn(node.name)
        try:
            doc_s = ast.get_docstring(node) or None
        except Exception:
            doc_s = None
        sym = Symbol(
            fqn=fqn,
            name=node.name,
            qualname=self._cur_qualname(),
            kind="class",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        self.class_stack.append(node.name)
        self.generic_visit(node)
        self.class_stack.pop()
        self.stack.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def _visit_func_like(self, node: Any) -> None:
        fqn = self._fqn(node.name)
        # Signature & returns
        sig_s, ret_s = None, None
        try:
            params = []
            for a in getattr(node, "args", None).args or []:
                nm = getattr(a, "arg", "")
                ann = getattr(a, "annotation", None)
                params.append(f"{nm}:{ast.unparse(ann)}" if ann is not None else nm)
            ret = getattr(node, "returns", None)
            ret_s = ast.unparse(ret) if ret is not None else None
            sig_s = f"({', '.join(params)})"


        except Exception:
            doc_s = None
        sym = Symbol(
            fqn=fqn,
            name=node.name,
            qualname=self._cur_qualname(),
            kind="class",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        self.class_stack.append(node.name)
        self.generic_visit(node)
        self.class_stack.pop()
        self.stack.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def _visit_func_like(self, node: Any) -> None:
        fqn = self._fqn(node.name)
        # Signature & returns
        sig_s, ret_s = None, None
        try:
            params = []
            for a in getattr(node, "args", None).args or []:
                nm = getattr(a, "arg", "")
                ann = getattr(a, "annotation", None)
                params.append(f"{nm}:{ast.unparse(ann)}" if ann is not None else nm)
            ret = getattr(node, "returns", None)
            ret_s = ast.unparse(ret) if ret is not None else None
            sig_s = f"({', '.join(params)})"
        except Exception:
            sig_s, ret_s = None, None
        try:


            fqn=fqn,
            name=node.name,
            qualname=self._cur_qualname(),
            kind="class",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        self.class_stack.append(node.name)
        self.generic_visit(node)
        self.class_stack.pop()
        self.stack.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def _visit_func_like(self, node: Any) -> None:
        fqn = self._fqn(node.name)
        # Signature & returns
        sig_s, ret_s = None, None
        try:
            params = []
            for a in getattr(node, "args", None).args or []:
                nm = getattr(a, "arg", "")
                ann = getattr(a, "annotation", None)
                params.append(f"{nm}:{ast.unparse(ann)}" if ann is not None else nm)
            ret = getattr(node, "returns", None)
            ret_s = ast.unparse(ret) if ret is not None else None
            sig_s = f"({', '.join(params)})"
        except Exception:
            sig_s, ret_s = None, None
        try:
            doc_s = ast.get_docstring(node) or None
        except Exception:
            doc_s = None


            kind="class",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        self.class_stack.append(node.name)
        self.generic_visit(node)
        self.class_stack.pop()
        self.stack.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:  # type: ignore[override]
        self._visit_func_like(node)

    def _visit_func_like(self, node: Any) -> None:
        fqn = self._fqn(node.name)
        # Signature & returns
        sig_s, ret_s = None, None
        try:
            params = []
            for a in getattr(node, "args", None).args or []:
                nm = getattr(a, "arg", "")
                ann = getattr(a, "annotation", None)
                params.append(f"{nm}:{ast.unparse(ann)}" if ann is not None else nm)
            ret = getattr(node, "returns", None)
            ret_s = ast.unparse(ret) if ret is not None else None
            sig_s = f"({', '.join(params)})"
        except Exception:
            sig_s, ret_s = None, None
        try:
            doc_s = ast.get_docstring(node) or None
        except Exception:
            doc_s = None
        sym = Symbol(
            fqn=fqn,
            name=node.name,
            qualname=self._cur_qualname(),
            kind="function",
            module=self.module,
            file=self.path,
            line=getattr(node, "lineno", 1),
            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
            signature=sig_s,
            returns=ret_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        # Traverse body to collect calls
        for sub in ast.walk(node):
            if isinstance(sub, ast.Call):
                callee_key = self._extract_callee_key(sub.func)
                if callee_key:
                    self.calls.append((fqn, callee_key))
        # Decorators as calls
        for dec in getattr(node, "decorator_list", []) or []:
            callee_key = self._extract_callee_key(dec)
            if callee_key:
                self.calls.append((fqn, callee_key))
        self.stack.pop()

    def visit_Assign(self, node: ast.Assign) -> Any:  # type: ignore[override]
        for t in getattr(node, "targets", []) or []:
            if isinstance(t, ast.Name):
                fqn = self._fqn(t.id)
                sym = Symbol(
                    fqn=fqn,
                    name=t.id,
                    qualname=self._cur_qualname(),
                    kind="variable",
                    module=self.module,
                    file=self.path,
                    line=getattr(node, "lineno", 1),
                    end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
                )
                self.symbols.append(sym)
        # capture __all__ = ["..."]
        try:
            names = []
            is_all = any(


            end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
            doc=doc_s,
            signature=sig_s,
            returns=ret_s,
        )
        self.symbols.append(sym)
        self.stack.append(node.name)
        # Traverse body to collect calls
        for sub in ast.walk(node):
            if isinstance(sub, ast.Call):
                callee_key = self._extract_callee_key(sub.func)
                if callee_key:
                    self.calls.append((fqn, callee_key))
        # Decorators as calls
        for dec in getattr(node, "decorator_list", []) or []:
            callee_key = self._extract_callee_key(dec)
            if callee_key:
                self.calls.append((fqn, callee_key))
        self.stack.pop()

    def visit_Assign(self, node: ast.Assign) -> Any:  # type: ignore[override]
        for t in getattr(node, "targets", []) or []:
            if isinstance(t, ast.Name):
                fqn = self._fqn(t.id)
                sym = Symbol(
                    fqn=fqn,
                    name=t.id,
                    qualname=self._cur_qualname(),
                    kind="variable",
                    module=self.module,
                    file=self.path,
                    line=getattr(node, "lineno", 1),
                    end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
                )
                self.symbols.append(sym)
        # capture __all__ = ["..."]
        try:
            names = []
            is_all = any(
                (isinstance(t, ast.Name) and t.id == "__all__") for t in node.targets
            )
            if is_all and isinstance(node.value, (ast.List, ast.Tuple)):
                for el in node.value.elts:
                    if isinstance(el, ast.Constant) and isinstance(el.value, str):
                        names.append(el.value)
            if names:
                self.exports.extend(names)
        except Exception:
            pass
        self.generic_visit(node)

    def _extract_callee_key(self, fn: ast.AST) -> Optional[str]:
        # simple name
        if isinstance(fn, ast.Name):
            return fn.id

        # super().method()
        if (
            isinstance(fn, ast.Attribute)
            and isinstance(fn.value, ast.Call)
            and isinstance(fn.value.func, ast.Name)
            and fn.value.func.id == "super"
        ):
            meth = fn.attr
            cur_cls = self._cur_class()
            if cur_cls:
                return f"{self.module}.{cur_cls}.{meth}"
            return meth

        # obj.attr chain


                    line=getattr(node, "lineno", 1),
                    end_line=getattr(node, "end_lineno", getattr(node, "lineno", 1)),
                )
                self.symbols.append(sym)
        # capture __all__ = ["..."]
        try:
            names = []
            is_all = any(
                (isinstance(t, ast.Name) and t.id == "__all__") for t in node.targets
            )
            if is_all and isinstance(node.value, (ast.List, ast.Tuple)):
                for el in node.value.elts:
                    if isinstance(el, ast.Constant) and isinstance(el.value, str):
                        names.append(el.value)
            if names:
                self.exports.extend(names)
        except Exception:
            pass
        self.generic_visit(node)

    def _extract_callee_key(self, fn: ast.AST) -> Optional[str]:
        # simple name
        if isinstance(fn, ast.Name):
            return fn.id

        # super().method()
        if (
            isinstance(fn, ast.Attribute)
            and isinstance(fn.value, ast.Call)
            and isinstance(fn.value.func, ast.Name)
            and fn.value.func.id == "super"
        ):
            meth = fn.attr
            cur_cls = self._cur_class()
            if cur_cls:
                return f"{self.module}.{cur_cls}.{meth}"
            return meth

        # obj.attr chain
        if isinstance(fn, ast.Attribute):
            parts: List[str] = []
            cur = fn
            while isinstance(cur, ast.Attribute):
                parts.append(cur.attr)
                cur = cur.value
            parts.reverse()

            if isinstance(cur, ast.Name):
                base = cur.id
                if base in ("self", "cls"):
                    cur_cls = self._cur_class()
                    if cur_cls and parts:
                        return f"{self.module}.{cur_cls}.{parts[-1]}"
                    return f"{self.module}.{cur_cls}" if cur_cls else parts[-1]
                if base in self.imports:
                    return f"{base}:{parts[-1]}" if parts else base
                return f"{self.module}.{base}.{parts[-1]}" if parts else base
        # getattr(module, "name") heuristic
        if (
            isinstance(fn, ast.Call)
            and isinstance(fn.func, ast.Name)
            and fn.func.id == "getattr"
            and fn.args
            and len(fn.args) >= 2
            and isinstance(fn.args[0], ast.Name)
            and isinstance(fn.args[1], ast.Constant)
            and isinstance(fn.args[1].value, str)
        ):
            base = fn.args[0].id
            name = fn.args[1].value
            if base in self.imports:
                return f"{self.imports[base]}.{name}"
            cur_cls = self._cur_class()
            if base in ("self", "cls") and cur_cls:
                return f"{self.module}.{cur_cls}.{name}"
            return f"{self.module}.{base}.{name}"
        # importlib.import_module("pkg.mod") heuristic
        if (
            isinstance(fn, ast.Call)
            and isinstance(fn.func, ast.Attribute)
            and isinstance(fn.func.value, ast.Name)
            and fn.func.value.id == "importlib"
            and fn.func.attr == "import_module"
            and fn.args
            and isinstance(fn.args[0], ast.Constant)
            and isinstance(fn.args[0].value, str)
        ):
            mod = str(fn.args[0].value)
            return mod
        return None

if __name__ == "__main__":
    _cli()


import types
from pathlib import Path
import sys

import torch

# Ensure project root is on sys.path so `tolbert` imports resolve consistently,
# even when tests are run from outside the repository root.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from tolbert import modeling

class _DummyEncoder(torch.nn.Module):
    """
    Tiny stand-in for a HuggingFace AutoModel encoder.

    It exposes:
      - .config.hidden_size


from pathlib import Path
import sys
from typing import List

import torch

# Ensure project root is on sys.path so we import the local `scripts` package.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts import build_repo_tree_of_life

class _DummySpan:
    def __init__(self, start_line: int, end_line: int) -> None:
        self.start_line = start_line
        self.end_line = end_line

class _DummyEntity:


import json
from pathlib import Path
import sys

import torch

# Ensure project root is on sys.path so local `tolbert` and `scripts` packages
# are importable regardless of the working directory.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from tolbert.data import TreeOfLifeDataset, collate_tree_of_life_batch
from scripts import eval_retrieval

def test_tree_of_life_dataset_and_collate_single_path(tmp_path: Path):
    """
    Verify that TreeOfLifeDataset:
      - reads node_path into per-level targets,
      - produces paths suitable for contrastive / path losses,


import csv
import json
from pathlib import Path
import sys
from typing import Dict

# Ensure project root is on sys.path so we import the local `scripts` package
# instead of any third-party package named `scripts`.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts import (
    build_wos_spans,
    build_researchhierarchy_spans,
    build_codehierarchy_spans,
    build_joint_code_paper_tol,
)

def test_build_wos_spans_helpers(tmp_path: Path):


import torch
from pathlib import Path
import sys

# Ensure project root is on sys.path so `tolbert` imports resolve consistently.
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from tolbert import losses
from tolbert import decoding

def test_tree_contrastive_loss_zero_when_no_shared_ancestors():
    """
    If no pairs share any non-root ancestor, the contrastive loss should be 0.
    """
    emb = torch.randn(3, 8)
    # All paths differ immediately after root.
    paths = [
        [0, 1],


tolbert/config.py

import json
from pathlib import Path
from typing import Any, Dict

def load_tolbert_config(path: str) -> Dict[str, Any]:
    """
    Load a simple YAML or JSON config file into a Python dict.

    - If the extension is .yaml or .yml, this function requires PyYAML.
    - If the extension is .json, it uses the standard library.
    """
    cfg_path = Path(path)
    if not cfg_path.exists():
        raise FileNotFoundError(f"Config file not found: {cfg_path}")

    suffix = cfg_path.suffix.lower()
    if suffix in {".yaml", ".yml"}:
        try:
            import yaml  # type: ignore
        except ImportError as e:


tolbert/modeling.py

from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

@dataclass
class TOLBERTConfig:
    """
    Minimal configuration for the TOLBERT encoder.

    This mirrors the sketch in `docs/api_reference.md` and is intended
    to be extended with any additional hyperparameters you need.
    """

    base_model_name: str
    level_sizes: Dict[int, int]
    proj_dim: int = 256


tolbert/losses.py

from typing import List, Sequence

import torch

def _compute_shared_depth(pi: Sequence[int], pj: Sequence[int]) -> int:
    """
    Depth of deepest shared label beyond the root.

    We treat index 0 in the path as the root. The shared depth is the
    number of consecutive matching labels starting from index 1.
    """
    max_len = min(len(pi), len(pj))
    depth = 0
    # Start from index 1 to skip the root.
    for idx in range(1, max_len):
        if pi[idx] != pj[idx]:
            break
        depth += 1
    return depth


"""
Minimal TOLBERT package skeleton.

This module exposes the core configuration and model classes so that
examples in `docs/usage.md` and `docs/api_reference.md` can import them.

The implementations here are intentionally lightweight and are meant
to be extended as you wire in real data and training code.
"""

from .modeling import TOLBERT, TOLBERTConfig

__all__ = ["TOLBERT", "TOLBERTConfig"]


tolbert/data.py

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Any

import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizerBase

@dataclass
class SpanRecord:
    text: str
    # Canonical single path used for per-level classification targets.
    node_path: Optional[List[int]]
    # Full raw record, which may include richer DAG-style fields such as
    # `node_paths` (multiple valid paths) in addition to `node_path`.
    raw: Dict[str, Any]


tolbert/decoding.py

from __future__ import annotations

from typing import Dict, List, Optional, Sequence

import torch

def greedy_hierarchical_decode(
    level_logits: Dict[str, torch.Tensor],
    parent_to_children: Dict[int, Dict[int, List[int]]],
    levels: Optional[Sequence[int]] = None,
) -> Dict[int, torch.Tensor]:
    """
    Greedy top-down decoding of a hierarchical path with child masking.

    This helper mirrors the inference-time procedure described in the TOLBERT
    paper: for each level, predictions are restricted to children of the
    previously chosen parent whenever a parent->children mapping is provided.

    Args:
        level_logits:


from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Tuple

# Basic identifiers
EntityId = str

@dataclass(frozen=True)
class Span:
    """Inclusive 1-based line span within an artifact."""

    start_line: int
    end_line: int

@dataclass
class Artifact:
    """


"""
Lightweight PDF preprocessing script.

Reads PDFs from /arxiv/pdfs/{year}/, extracts structured tokens via P0's tokenizer,
and writes JSONL chunks (section/equation/figure/table/text) under exports/pdfs_structured/.

Usage:
  # Specific year range (inclusive)
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --years 2018 2020 --max-files 1000

  # All years and all files (may be large)
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0

  # Use a VLM (e.g., Qwen/Qwen3-VL-2B-Instruct) for OCR
  PYTHONPATH=.. python -m models.scripts.preprocess_pdfs --max-files 0 --qwen-model Qwen/Qwen3-VL-2B-Instruct

Supports resumable runs: existing shards in the output directory are scanned and
their pdf_path entries are skipped, so re-running continues where it left off.
"""

import argparse


"""
Simple cross-domain zero-shot evaluation script for TOLBERT.

This wraps the hierarchical classification evaluator to support the scenario:

  - Train TOLBERT on domain A (e.g., CodeHierarchy).
  - Evaluate the same checkpoint on domain B (e.g., WOS) without further training.

Usage:

  # Evaluate a code-trained model on WOS spans
  python -m scripts.eval_zero_shot_cross_domain \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint checkpoints/codehierarchy/tolbert_epoch5.pt \\
      --target-config configs/wos_example.yaml \\
      --target-spans data/wos/spans_test.jsonl

The script uses:
  - base_model_name and model head structure from the *source* config
    (the one used for training),
  - but evaluation data and tokenizer settings (e.g., max_length) from


"""
Build JSONL spans and ontology metadata for a WOS-style hierarchical text dataset.

This script is a reference implementation of the ResearchHierarchy construction
described in the TOLBERT paper. It does **not** ship any WOS data; instead, it
expects you to provide a CSV file with one row per document and three levels of
labels.

Expected input CSV schema (you can adapt this to your variant as needed):

    text_col:   column containing the text to encode (e.g., abstract)
    l1_col:     top-level field (e.g., "Computer Science")
    l2_col:     second-level subfield
    l3_col:     third-level discipline

By default we assume:

    --text-col text
    --l1-col   level1
    --l2-col   level2
    --l3-col   level3


from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str


"""
Build ResearchHierarchy-style JSONL spans and simple ontology metadata for WOS / arXiv papers.

This is a reference implementation of the dataset construction described for the
ResearchHierarchy benchmark in the TOLBERT paper. It expects you to provide a local
metadata file that assigns each paper to a 3-level research taxonomy and (optionally)
its backing PDF path.

The script is intentionally lightweight and mirrors the structure of
`build_codehierarchy_spans.py`:

Input
=====
- --metadata_file:
    A CSV or JSONL file with, at minimum, the columns:
      doc_id, field, subfield, discipline, text

    Optionally it may also include:
      pdf_path: absolute or repo-relative path to the PDF for this paper
      source:  short tag like "wos" or "arxiv" identifying the corpus


"""
Agent / LLM stack skeleton: PR -> paper recommendations via TOLBERT.

This script is a high-level wiring of the "worked scenario" from `docs/usage.md`:
  1. Ingest a pull request (PR) as text spans.
  2. Encode spans with TOLBERT.
  3. Locate the PR within the Tree-of-Life (domain/subdomain).
  4. Retrieve relevant paper spans from a pre-built index.
  5. Use an LLM to summarize / explain the recommendations.

You are expected to plug in:
  - real PR ingestion,
  - a proper paper-span index (vector DB, search service, etc.),
  - and your own LLM backend in `call_llm`.
"""

import argparse
from pathlib import Path
from typing import Any, Dict, List, Tuple

import torch


"""
Build a joint code+paper Tree-of-Life (ToL) from existing per-domain spans.

This script is the missing glue described in the paper/docs: it takes
separate code and paper span files (each with its own local ontology and
`node_path` ids) and produces:

- A single, unified `nodes.jsonl` Tree-of-Life with:
    level 0: root
    level 1: domain nodes (e.g., "Code", "Papers")
    level 2+: remapped nodes from the original per-domain ontologies
- Rewritten span files for code and papers where `node_path` has been
  updated to refer to the *joint* node ID space.
- A `level_sizes` JSON helper compatible with `TOLBERTConfig.level_sizes`.

Design notes
============

This script deliberately **does not attempt to infer a semantic mapping**
between code categories (e.g., "Web", "ML", "Systems") and paper fields
("Computer Science", "Physics", ...). Doing that well would require


"""
Evaluate a flat (leaf-level) baseline classifier on a labeled spans_file.

This computes standard leaf-level accuracy for the "BERT-flat" style baselines
trained with scripts/train_flat_baseline.py. It does NOT reconstruct full
hierarchical paths; it simply evaluates how well the model predicts the leaf
node id (last element of node_path).

Example:

  python -m scripts.eval_flat_baseline \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint-dir checkpoints/codehierarchy_bert_flat \\
      --spans-file /data/tolbert/data/codehierarchy/spans_test.jsonl
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path


"""
Minimal training skeleton for TOLBERT **with curriculum and multi-domain support**.

This script wires together:
  - config loading
  - tokenizer + dataset(s)
  - model + optimizer
  - a training loop with optional curriculum over hierarchical / path / contrastive losses

It assumes you have prepared one or more `spans_file`(s) as described in `docs/tree_of_life.md`.
"""

import argparse
from pathlib import Path
from typing import Any, Dict, Optional

import torch
from torch.utils.data import DataLoader, ConcatDataset
from transformers import AutoTokenizer
import os


from __future__ import annotations

import os
import re
import hashlib
from typing import Iterable, List, Tuple, Dict, Optional, Set, Any

from modules.program_graph import ProgramGraph, Entity, Edge, Artifact, Span, ResolvedAnchor, EntityId

def program_id_for_repo(repo_root: str) -> str:
    base = os.path.basename(os.path.abspath(repo_root)) or "repo"
    return base

def artifact_uri(program_id: str, rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    return f"program://{program_id}/artifact/{rel}"

def parse_program_uri(uri: str) -> Tuple[str, str, str, Optional[Tuple[int, int]]]:


"""
Evaluate a trained TOLBERT checkpoint on a labeled spans_file.

This script implements the hierarchical classification metrics discussed
in the TOLBERT paper:

  - per-level accuracy
  - per-level micro F1 (for single-label multiclass this equals accuracy)
  - path accuracy (all supervised levels correct for a span)

Usage (CodeHierarchy example):

  python -m scripts.eval_hierarchical_classification \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint checkpoints/codehierarchy/tolbert_epoch5.pt \\
      --spans-file data/codehierarchy/spans_test.jsonl
"""

from __future__ import annotations

import argparse


"""
Build a simple Tree-of-Life taxonomy and span records from a single repository
using the language-agnostic RepoGraph.

This is a concrete, graph-driven implementation of the code-side part of
`docs/tree_of_life.md`:

  - Nodes:
      level 0: root
      level 1: language nodes (Python, Java, C/C++, Go, JS/TS, ...)
      level 2: repo node
      level 3: top-level directory nodes (subtrees within the repo)
      level 4: file nodes
      level 5: symbol nodes (functions / methods / classes), when available

  - Edges:
      parent_id / child_id edges forming a tree over the above nodes.

  - Spans:
      one span per file, with a node_path from root → language → repo →
      top-level directory → file.


"""
Build CodeHierarchy-style JSONL spans and simple ontology metadata.

This is a *reference* implementation of the dataset construction described
in the TOLBERT paper for the CodeHierarchy benchmark. It does not fetch or
ship any GitHub data; instead, it expects you to point it at a directory of
local repositories plus a small metadata file that assigns each repo to a
language and coarse category.

Input:
  - --repos_root:
      Directory that contains one subdirectory per repository, e.g.
        repos_root/
          repo1/
            ...
          repo2/
            ...
  - --metadata_file:
      A CSV or JSONL file with, at minimum, the fields:
        repo_name, language, category
      where:


"""
Build ArXiv-CLS–style JSONL spans and simple ontology metadata.

This script is a small helper for the ArXiv-CLS variant described in the
TOLBERT paper. It mirrors the style of `build_wos_spans.py` / `build_codehierarchy_spans.py`,
but targets the 2-level arXiv category hierarchy.

The ArXiv-CLS setup in the paper uses:

  - Documents: recent arXiv CS papers (e.g., title + abstract text).
  - Labels: a two-level arXiv category taxonomy, e.g.:
        domain  = "cs"
        category = "cs.LG"

You are expected to provide a metadata file (CSV or JSONL) with at least:

  - doc_id:   unique paper identifier (e.g., arXiv ID)
  - domain:   top-level arXiv area (e.g., "cs", "math")
  - category: fine-grained arXiv category (e.g., "cs.LG", "cs.CL")
  - text:     text span to encode (e.g., title + abstract)


"""
Evaluate embeddings on retrieval-style tasks.

This script computes:
  - MRR (Mean Reciprocal Rank)
  - Precision@K

for a simple setup where:
  - The index consists of spans from one spans_file (e.g., code files).
  - The queries are spans from another spans_file (e.g., paper abstracts).
  - "Relevance" is defined by sharing one or more levels in the hierarchy.

By default we use "share at least level 2" as the relevance criterion, but
this can be adjusted via flags.

Usage (Paper2Code-style example):

  python -m scripts.eval_retrieval \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint checkpoints/tolbert_epoch5.pt \\
      --index-spans data/codehierarchy/spans_code.jsonl \\


"""
Local scripts package for TOLBERT utilities (dataset builders, training, eval).

This file exists so that tests and external tools can reliably import modules
like `scripts.build_wos_spans` without accidentally resolving to any unrelated
third-party `scripts` package on the Python path.
"""


"""
Train a flat (leaf-level) baseline classifier on spans_file.

This script is intended to reproduce the "BERT-flat" style baselines from
Section 4 using standard HuggingFace models (e.g. BERT, SciBERT, CodeBERT,
ModernBERT) on the same JSONL spans files used for TOLBERT.

Key characteristics:
  - Single-label, flat classification on the *leaf* node of each path
    (last element of `node_path`).
  - Uses AutoModelForSequenceClassification with a single softmax head.
  - Ignores intermediate hierarchy levels entirely.

Example (CodeHierarchy, BERT-base flat):

  python -m scripts.train_flat_baseline \\
      --config configs/codehierarchy_example.yaml \\
      --output-dir checkpoints/codehierarchy_bert_flat \\
      --base-model-name bert-base-uncased

Example (WOS, SciBERT flat):


"""
Simple retrieval sandbox for TOLBERT.

This script:
  - loads a trained TOLBERT checkpoint and tokenizer,
  - encodes spans from a `spans_file`,
  - builds an in-memory index of embeddings,
  - lets you run ad-hoc queries over that index from the command line.

It is intended for small-scale experimentation and debugging, not production.
"""

import argparse
from pathlib import Path
from typing import Any, Dict, List, Tuple

import torch
from torch.nn import functional as F
from transformers import AutoTokenizer
import os


import os
import ast
import re
import json
import time  # noqa: F401
from dataclasses import dataclass, field
from typing import Any, Dict, List, Tuple, Optional
try:
    import pathspec  # type: ignore
except Exception:  # pragma: no cover
    pathspec = None  # type: ignore

@dataclass
class Symbol:
    fqn: str
    name: str
    qualname: str
    kind: str  # module|class|function|variable
    module: str
    file: str