code-search-cli 0.1.0__py3-none-any.whl

code_retrieval/index_store.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import sqlite3
+import time
+from collections import Counter
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from .models import Section, SourceFile
+from .query import document_terms
+from .repository import RepositoryReader, SourceSignature
+from .structure import sections_for
+
+
+SCHEMA_VERSION = "1"
+
+
+@dataclass(slots=True)
+class IndexSnapshot:
+    files: list[SourceFile]
+    document_terms: dict[str, Counter[str]]
+    sections: dict[str, list[Section]]
+
+
+@dataclass(slots=True)
+class IndexRefresh:
+    index_path: str
+    source_files: int
+    added: int
+    updated: int
+    removed: int
+    unchanged: int
+    elapsed_ms: float
+
+    def to_dict(self) -> dict[str, str | int | float]:
+        return asdict(self)
+
+
+@dataclass(slots=True)
+class IndexStatus:
+    index_path: str
+    exists: bool
+    fresh: bool
+    source_files: int
+    added: int
+    updated: int
+    removed: int
+    revision: str | None
+
+    def to_dict(self) -> dict[str, str | int | bool | None]:
+        return asdict(self)
+
+
+def cache_root() -> Path:
+    override = os.environ.get("CODE_SEARCH_CACHE_HOME")
+    if override:
+        return Path(override).expanduser()
+    xdg = os.environ.get("XDG_CACHE_HOME")
+    return (Path(xdg).expanduser() if xdg else Path.home() / ".cache") / "code-search"
+
+
+class PersistentIndex:
+    def __init__(self, reader: RepositoryReader):
+        self.reader = reader
+        identity = f"{reader.root}\0{reader.ref or 'WORKTREE'}".encode()
+        self.path = cache_root() / f"{hashlib.sha256(identity).hexdigest()[:24]}.sqlite3"
+
+    def refresh(self) -> IndexRefresh:
+        started = time.perf_counter()
+        self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+        if not self.path.exists():
+            self.path.touch(mode=0o600)
+        with self._connect() as connection:
+            self._initialize(connection)
+            if self.reader.ref:
+                counts = self._refresh_revision(connection)
+            else:
+                counts = self._refresh_worktree(connection)
+            source_files = int(connection.execute("SELECT COUNT(*) FROM files").fetchone()[0])
+        return IndexRefresh(
+            str(self.path),
+            source_files,
+            counts["added"],
+            counts["updated"],
+            counts["removed"],
+            counts["unchanged"],
+            round((time.perf_counter() - started) * 1000, 2),
+        )
+
+    def load(self) -> IndexSnapshot:
+        files = []
+        terms = {}
+        sections = {}
+        with self._connect() as connection:
+            self._initialize(connection)
+            rows = connection.execute(
+                "SELECT path, content, terms, sections FROM files ORDER BY path"
+            )
+            for path, content, raw_terms, raw_sections in rows:
+                source = SourceFile(path, content)
+                files.append(source)
+                terms[path] = Counter(json.loads(raw_terms))
+                lines = content.splitlines()
+                sections[path] = [
+                    Section(
+                        path,
+                        item["name"],
+                        item["kind"],
+                        item["start_line"],
+                        item["end_line"],
+                        "\n".join(lines[item["start_line"] - 1 : item["end_line"]]),
+                    )
+                    for item in json.loads(raw_sections)
+                ]
+        return IndexSnapshot(files, terms, sections)
+
+    def status(self) -> IndexStatus:
+        if not self.path.exists():
+            return IndexStatus(str(self.path), False, False, 0, 0, 0, 0, self.reader.resolved_revision())
+        with self._connect() as connection:
+            self._initialize(connection)
+            source_files = int(connection.execute("SELECT COUNT(*) FROM files").fetchone()[0])
+            if self.reader.ref:
+                revision = self.reader.resolved_revision()
+                indexed_revision = self._meta(connection, "revision")
+                fresh = self._meta(connection, "initialized") == "1" and revision == indexed_revision
+                return IndexStatus(
+                    str(self.path), True, fresh, source_files,
+                    0 if fresh else source_files, 0, 0 if fresh else source_files, revision,
+                )
+            current = {item.path: item for item in self.reader.source_signatures()}
+            stored = self._stored_signatures(connection)
+            added = len(current.keys() - stored.keys())
+            removed = len(stored.keys() - current.keys())
+            updated = sum(
+                current[path].size != stored[path][0] or current[path].mtime_ns != stored[path][1]
+                for path in current.keys() & stored.keys()
+            )
+            return IndexStatus(
+                str(self.path), True, not (added or updated or removed), source_files,
+                added, updated, removed, None,
+            )
+
+    def clean(self) -> bool:
+        removed = False
+        for path in (self.path, Path(f"{self.path}-wal"), Path(f"{self.path}-shm")):
+            if path.exists():
+                path.unlink()
+                removed = True
+        return removed
+
+    def _refresh_worktree(self, connection: sqlite3.Connection) -> dict[str, int]:
+        current = {item.path: item for item in self.reader.source_signatures()}
+        stored = self._stored_signatures(connection)
+        added_paths = current.keys() - stored.keys()
+        removed_paths = stored.keys() - current.keys()
+        updated_paths = {
+            path for path in current.keys() & stored.keys()
+            if current[path].size != stored[path][0] or current[path].mtime_ns != stored[path][1]
+        }
+        for path in removed_paths:
+            connection.execute("DELETE FROM files WHERE path = ?", (path,))
+        indexed_added = 0
+        indexed_updated = 0
+        for path in sorted(added_paths | updated_paths):
+            try:
+                source = self.reader.read_file(path)
+            except OSError:
+                connection.execute("DELETE FROM files WHERE path = ?", (path,))
+                continue
+            if "\x00" in source.content:
+                connection.execute("DELETE FROM files WHERE path = ?", (path,))
+                continue
+            signature = current[path]
+            self._store_source(connection, source, signature)
+            if path in added_paths:
+                indexed_added += 1
+            else:
+                indexed_updated += 1
+        self._set_meta(connection, "revision", "")
+        self._set_meta(connection, "initialized", "1")
+        return {
+            "added": indexed_added,
+            "updated": indexed_updated,
+            "removed": len(removed_paths),
+            "unchanged": len(current) - len(added_paths) - len(updated_paths),
+        }
+
+    def _refresh_revision(self, connection: sqlite3.Connection) -> dict[str, int]:
+        revision = self.reader.resolved_revision()
+        previous = self._meta(connection, "revision")
+        initialized = self._meta(connection, "initialized") == "1"
+        previous_count = int(connection.execute("SELECT COUNT(*) FROM files").fetchone()[0])
+        if initialized and revision == previous:
+            return {"added": 0, "updated": 0, "removed": 0, "unchanged": previous_count}
+        connection.execute("DELETE FROM files")
+        added = 0
+        for source in self.reader.iter_source_files():
+            self._store_source(connection, source, SourceSignature(source.path, len(source.content), 0))
+            added += 1
+        self._set_meta(connection, "revision", revision or "")
+        self._set_meta(connection, "initialized", "1")
+        return {"added": added, "updated": 0, "removed": previous_count, "unchanged": 0}
+
+    def _store_source(
+        self,
+        connection: sqlite3.Connection,
+        source: SourceFile,
+        signature: SourceSignature,
+    ) -> None:
+        section_metadata = [
+            {
+                "name": item.name,
+                "kind": item.kind,
+                "start_line": item.start_line,
+                "end_line": item.end_line,
+            }
+            for item in sections_for(source)
+        ]
+        connection.execute(
+            """
+            INSERT INTO files(path, size, mtime_ns, content, terms, sections)
+            VALUES (?, ?, ?, ?, ?, ?)
+            ON CONFLICT(path) DO UPDATE SET
+                size=excluded.size,
+                mtime_ns=excluded.mtime_ns,
+                content=excluded.content,
+                terms=excluded.terms,
+                sections=excluded.sections
+            """,
+            (
+                source.path,
+                signature.size,
+                signature.mtime_ns,
+                source.content,
+                json.dumps(document_terms(source.content), separators=(",", ":")),
+                json.dumps(section_metadata, separators=(",", ":")),
+            ),
+        )
+
+    def _stored_signatures(self, connection: sqlite3.Connection) -> dict[str, tuple[int, int]]:
+        return {
+            path: (int(size), int(mtime_ns))
+            for path, size, mtime_ns in connection.execute("SELECT path, size, mtime_ns FROM files")
+        }
+
+    def _connect(self) -> sqlite3.Connection:
+        return sqlite3.connect(self.path, timeout=30)
+
+    def _initialize(self, connection: sqlite3.Connection) -> None:
+        connection.execute("CREATE TABLE IF NOT EXISTS metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL)")
+        version = self._meta(connection, "schema_version")
+        if version and version != SCHEMA_VERSION:
+            connection.execute("DROP TABLE IF EXISTS files")
+            connection.execute("DELETE FROM metadata")
+        connection.execute(
+            """
+            CREATE TABLE IF NOT EXISTS files(
+                path TEXT PRIMARY KEY,
+                size INTEGER NOT NULL,
+                mtime_ns INTEGER NOT NULL,
+                content TEXT NOT NULL,
+                terms TEXT NOT NULL,
+                sections TEXT NOT NULL
+            )
+            """
+        )
+        self._set_meta(connection, "schema_version", SCHEMA_VERSION)
+        self._set_meta(connection, "repository", str(self.reader.root))
+        self._set_meta(connection, "ref", self.reader.ref or "")
+
+    @staticmethod
+    def _meta(connection: sqlite3.Connection, key: str) -> str | None:
+        row = connection.execute("SELECT value FROM metadata WHERE key = ?", (key,)).fetchone()
+        return str(row[0]) if row else None
+
+    @staticmethod
+    def _set_meta(connection: sqlite3.Connection, key: str, value: str) -> None:
+        connection.execute(
+            "INSERT INTO metadata(key, value) VALUES (?, ?) "
+            "ON CONFLICT(key) DO UPDATE SET value=excluded.value",
+            (key, value),
+        )

code_retrieval/installer.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import os
+from importlib.resources import files
+from pathlib import Path
+from typing import Literal
+
+
+InstallAction = Literal["created", "updated", "unchanged"]
+SKILL_FILES = (Path("SKILL.md"), Path("agents/openai.yaml"))
+
+
+def codex_skills_root() -> Path:
+    codex_home = os.environ.get("CODEX_HOME")
+    return (Path(codex_home).expanduser() if codex_home else Path.home() / ".codex") / "skills"
+
+
+def install_codex_skill(*, force: bool = False) -> tuple[Path, InstallAction]:
+    target = codex_skills_root() / "code-search"
+    source = files("code_retrieval").joinpath("resources/code-search")
+    existed = target.exists()
+    changed = False
+
+    for relative in SKILL_FILES:
+        destination = target / relative
+        content = source.joinpath(relative.as_posix()).read_text(encoding="utf-8")
+        if destination.exists():
+            existing = destination.read_text(encoding="utf-8")
+            if existing == content:
+                continue
+            if not force:
+                raise FileExistsError(
+                    f"refusing to overwrite modified skill file: {destination}; rerun with --force"
+                )
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        destination.write_text(content, encoding="utf-8")
+        changed = True
+
+    action: InstallAction = "unchanged" if not changed else ("updated" if existed else "created")
+    return target, action
+
+
+def uninstall_codex_skill(*, force: bool = False) -> tuple[Path, Literal["removed", "not-found"]]:
+    target = codex_skills_root() / "code-search"
+    if not target.exists():
+        return target, "not-found"
+
+    source = files("code_retrieval").joinpath("resources/code-search")
+    for relative in SKILL_FILES:
+        destination = target / relative
+        if not destination.exists():
+            continue
+        expected = source.joinpath(relative.as_posix()).read_text(encoding="utf-8")
+        if destination.read_text(encoding="utf-8") != expected and not force:
+            raise FileExistsError(
+                f"refusing to remove modified skill file: {destination}; rerun with --force"
+            )
+        destination.unlink()
+
+    agents = target / "agents"
+    if agents.exists() and not any(agents.iterdir()):
+        agents.rmdir()
+    if target.exists() and not any(target.iterdir()):
+        target.rmdir()
+    return target, "removed"

code_retrieval/models.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+
+@dataclass(slots=True)
+class SourceFile:
+    path: str
+    content: str
+
+
+@dataclass(slots=True)
+class Section:
+    path: str
+    name: str
+    kind: str
+    start_line: int
+    end_line: int
+    content: str
+
+
+@dataclass(slots=True)
+class Evidence:
+    path: str
+    start_line: int
+    end_line: int
+    kind: str
+    name: str
+    score: float
+    reasons: list[str]
+    content: str
+    estimated_tokens: int
+
+    def key(self) -> tuple[str, int, int]:
+        return self.path, self.start_line, self.end_line
+
+
+@dataclass(slots=True)
+class RetrievalResult:
+    task: str
+    repository: str
+    ref: str | None
+    used_tokens: int
+    query_terms: list[str]
+    evidence: list[Evidence]
+    stats: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def files(self) -> list[str]:
+        return list(dict.fromkeys(item.path for item in self.evidence))
+
+    def to_dict(self) -> dict[str, Any]:
+        data = asdict(self)
+        data["files"] = self.files
+        return data
+
+    def to_markdown(self) -> str:
+        ref = self.ref or "working tree"
+        lines = [
+            "# Retrieval evidence",
+            "",
+            f"- Task: {self.task}",
+            f"- Repository: `{self.repository}`",
+            f"- Revision: `{ref}`",
+            f"- Files: {len(self.files)}; evidence spans: {len(self.evidence)}; "
+            f"estimated payload: {self.used_tokens} tokens",
+            "",
+        ]
+        for index, item in enumerate(self.evidence, 1):
+            suffix = _language_for(item.path)
+            lines.extend(
+                [
+                    f"## {index}. `{item.path}:{item.start_line}-{item.end_line}`",
+                    "",
+                    f"Score `{item.score:.2f}` · {item.kind} `{item.name}` · "
+                    + "; ".join(item.reasons),
+                    "",
+                    f"```{suffix}",
+                    item.content.rstrip(),
+                    "```",
+                    "",
+                ]
+            )
+        return "\n".join(lines)
+
+
+def _language_for(path: str) -> str:
+    suffix = path.rsplit(".", 1)[-1].lower() if "." in path else "text"
+    return {"py": "python", "js": "javascript", "ts": "typescript"}.get(suffix, suffix)

code_retrieval/query.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import re
+from collections import Counter
+
+
+WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9_]*|[\u4e00-\u9fff]+")
+CAMEL_RE = re.compile(r"[A-Z]+(?=[A-Z][a-z]|\d|$)|[A-Z]?[a-z]+|\d+")
+
+STOP_WORDS = {
+    "a", "an", "and", "are", "be", "for", "from", "in", "is", "it", "of", "on", "or",
+    "so", "that", "the", "this", "to", "with", "not", "fix", "resolve", "detect", "issue", "problem",
+    "missing", "missed", "change", "changes", "code", "configuration",
+}
+
+# Generic engineering vocabulary, not repository- or competitor-specific rules.
+ALIASES: dict[str, tuple[str, ...]] = {
+    "nullpointerexception": ("null", "nullable", "optional"),
+    "npe": ("null", "nullable"),
+    "scheduled": ("schedule", "scheduler", "cron", "trigger"),
+    "schedule": ("scheduler", "cron", "trigger"),
+    "configuration": ("config", "properties", "settings"),
+    "config": ("configuration", "properties", "settings"),
+    "callback": ("listener", "subscriber", "event", "notify", "publish", "handler"),
+    "utf8": ("utf", "charset", "encoding", "encode"),
+    "md5": ("digest", "checksum", "hash"),
+    "dependency": ("depends", "dependson", "bean"),
+    "injection": ("autowired", "inject", "bean", "component"),
+    "order": ("dependson", "postconstruct", "initialize", "initialization", "init"),
+    "startup": ("postconstruct", "initialize", "initialization", "init", "bean"),
+    "event": ("notify", "publish", "subscriber", "listener"),
+    "events": ("event", "notify", "publish", "subscriber", "listener"),
+}
+
+
+def split_identifier(value: str) -> list[str]:
+    parts: list[str] = []
+    for raw in WORD_RE.findall(value.replace("_", " ")):
+        if raw.isascii():
+            camel = CAMEL_RE.findall(raw)
+            parts.extend(piece.lower() for piece in camel if piece)
+            joined = raw.lower()
+            if joined not in parts:
+                parts.append(joined)
+        else:
+            parts.append(raw)
+    return parts
+
+
+def document_terms(value: str) -> Counter[str]:
+    return Counter(split_identifier(value))
+
+
+def query_terms(task: str) -> tuple[dict[str, float], list[str]]:
+    raw = split_identifier(task)
+    original = list(dict.fromkeys(term for term in raw if term not in STOP_WORDS and len(term) > 1))
+    weighted = {term: 3.0 for term in original}
+    for term in raw:
+        for alias in ALIASES.get(term, ()):
+            weighted.setdefault(alias, 1.0)
+    return weighted, original

code_retrieval/repository.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterator
+
+from .models import SourceFile
+
+
+SOURCE_SUFFIXES = {
+    ".c", ".cc", ".cpp", ".cs", ".go", ".h", ".hpp", ".java", ".js", ".jsx", ".kt",
+    ".kts", ".php", ".py", ".rb", ".rs", ".scala", ".sh", ".sql", ".swift", ".ts", ".tsx",
+    ".vue", ".xml", ".yaml", ".yml", ".toml", ".properties", ".gradle",
+}
+SKIP_DIRS = {
+    ".git", ".idea", ".mypy_cache", ".pytest_cache", ".tox", ".venv", "build", "coverage",
+    "dist", "node_modules", "target", "vendor", "venv",
+}
+MAX_FILE_BYTES = 1_000_000
+
+
+@dataclass(frozen=True, slots=True)
+class SourceSignature:
+    path: str
+    size: int
+    mtime_ns: int
+
+
+class RepositoryReader:
+    def __init__(self, root: str | Path, ref: str | None = None):
+        self.root = Path(root).expanduser().resolve()
+        self.ref = ref
+        if not self.root.is_dir():
+            raise ValueError(f"repository is not a directory: {self.root}")
+        if ref:
+            self._verify_ref(ref)
+
+    def iter_source_files(self) -> Iterator[SourceFile]:
+        if self.ref:
+            yield from self._iter_git_ref()
+        else:
+            yield from self._iter_worktree()
+
+    def source_signatures(self) -> list[SourceSignature]:
+        if self.ref:
+            raise ValueError("source signatures are only available for a working tree")
+        signatures = []
+        for path, relative in self._worktree_paths():
+            try:
+                stat = path.stat()
+            except OSError:
+                continue
+            if stat.st_size <= MAX_FILE_BYTES:
+                signatures.append(SourceSignature(relative, stat.st_size, stat.st_mtime_ns))
+        return signatures
+
+    def resolved_revision(self) -> str | None:
+        if not self.ref:
+            return None
+        return subprocess.run(
+            ["git", "rev-parse", "--verify", f"{self.ref}^{{commit}}"],
+            cwd=self.root,
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout.strip()
+
+    def read_file(self, path: str) -> SourceFile:
+        if not _is_source(path):
+            raise ValueError(f"unsupported source file: {path}")
+        if self.ref:
+            result = subprocess.run(
+                ["git", "show", f"{self.ref}:{path}"], cwd=self.root, capture_output=True, check=True
+            )
+            return SourceFile(path, _decode(result.stdout))
+        resolved = (self.root / path).resolve()
+        if not resolved.is_relative_to(self.root):
+            raise ValueError("path escapes repository")
+        return SourceFile(path, _decode(resolved.read_bytes()))
+
+    def _iter_worktree(self) -> Iterator[SourceFile]:
+        for path, relative in self._worktree_paths():
+            try:
+                if path.stat().st_size > MAX_FILE_BYTES:
+                    continue
+                data = path.read_bytes()
+            except OSError:
+                continue
+            if b"\x00" not in data:
+                yield SourceFile(relative, _decode(data))
+
+    def _worktree_paths(self) -> Iterator[tuple[Path, str]]:
+        for current, dirs, files in os.walk(self.root):
+            dirs[:] = sorted(name for name in dirs if name not in SKIP_DIRS)
+            for name in sorted(files):
+                path = Path(current, name)
+                relative = path.relative_to(self.root).as_posix()
+                if not _is_source(relative):
+                    continue
+                try:
+                    if not path.resolve().is_relative_to(self.root):
+                        continue
+                except OSError:
+                    continue
+                yield path, relative
+
+    def _iter_git_ref(self) -> Iterator[SourceFile]:
+        listing = subprocess.run(
+            ["git", "ls-tree", "-r", "-z", "--name-only", self.ref],
+            cwd=self.root,
+            capture_output=True,
+            check=True,
+        ).stdout
+        paths = [path.decode("utf-8", "surrogateescape") for path in listing.split(b"\x00") if path]
+        paths = [path for path in paths if _is_source(path) and not _has_skipped_dir(path)]
+        process = subprocess.Popen(
+            ["git", "cat-file", "--batch"],
+            cwd=self.root,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+        )
+        assert process.stdin is not None and process.stdout is not None
+        try:
+            for path in paths:
+                process.stdin.write(f"{self.ref}:{path}\n".encode())
+                process.stdin.flush()
+                header = process.stdout.readline().decode().strip().split()
+                if len(header) != 3 or header[1] != "blob":
+                    continue
+                size = int(header[2])
+                data = process.stdout.read(size)
+                process.stdout.read(1)
+                if size <= MAX_FILE_BYTES and b"\x00" not in data:
+                    yield SourceFile(path, _decode(data))
+        finally:
+            process.stdin.close()
+            process.stdout.close()
+            process.wait(timeout=10)
+
+    def _verify_ref(self, ref: str) -> None:
+        result = subprocess.run(
+            ["git", "rev-parse", "--verify", f"{ref}^{{commit}}"],
+            cwd=self.root,
+            capture_output=True,
+        )
+        if result.returncode:
+            raise ValueError(f"unknown git revision: {ref}")
+
+
+def _is_source(path: str) -> bool:
+    return Path(path).suffix.lower() in SOURCE_SUFFIXES
+
+
+def _has_skipped_dir(path: str) -> bool:
+    return any(part in SKIP_DIRS for part in Path(path).parts[:-1])
+
+
+def _decode(data: bytes) -> str:
+    return data.decode("utf-8", errors="replace")