diffcontext 0.5.1__py3-none-any.whl

diffcontext/lexical.py

@@ -0,0 +1,108 @@
+"""
+lexical.py — BM25 lexical retrieval signal.
+
+Pure-stdlib BM25Okapi over symbol source code, used as the lexical leg of
+hybrid retrieval. The eval_v2 benchmark showed the call graph alone loses
+to full-code BM25 on most measures, while the graph+BM25+same-file blend
+beats every individual signal on 4/5 repos (see benchmarks/EVAL_V2_REPORT.md);
+this module is that lexical leg, with no third-party dependency.
+
+The math replicates rank_bm25's BM25Okapi (k1=1.5, b=0.75, and negative-idf
+flooring at epsilon * average_idf) so product scores match the benchmarked
+implementation. Scoring uses an inverted index, so a query only touches
+documents that share at least one term with it — indexing a ~9k-symbol repo
+takes ~1s and a query a few milliseconds.
+"""
+
+import math
+import re
+from collections import Counter, defaultdict
+from typing import Dict, List
+
+from .models import Symbol
+
+K1 = 1.5
+B = 0.75
+EPSILON = 0.25   # floor for negative idf, as a fraction of average idf
+
+_TOKEN_RE = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
+
+
+def tokenize(code: str) -> List[str]:
+    """Identifier tokens, lowercased, single-char tokens dropped."""
+    return [t.lower() for t in _TOKEN_RE.findall(code) if len(t) > 1]
+
+
+class LexicalIndex:
+    """BM25 index over every symbol's source code."""
+
+    def __init__(self, symbols: Dict[str, Symbol]):
+        self.ids: List[str] = list(symbols.keys())
+        n = len(self.ids)
+
+        doc_len: List[int] = []
+        term_doc_freq: Counter = Counter()          # term -> #docs containing it
+        postings: Dict[str, List] = defaultdict(list)  # term -> [(doc_idx, tf)]
+
+        for i, sid in enumerate(self.ids):
+            tf = Counter(tokenize(symbols[sid].code))
+            doc_len.append(sum(tf.values()))
+            for term, count in tf.items():
+                term_doc_freq[term] += 1
+                postings[term].append((i, count))
+
+        self.doc_len = doc_len
+        self.avgdl = sum(doc_len) / n if n else 0.0
+        self.postings = postings
+
+        # idf with rank_bm25's negative-idf flooring
+        idf: Dict[str, float] = {}
+        idf_sum, negatives = 0.0, []
+        for term, df in term_doc_freq.items():
+            v = math.log(n - df + 0.5) - math.log(df + 0.5)
+            idf[term] = v
+            idf_sum += v
+            if v < 0:
+                negatives.append(term)
+        if idf:
+            # rank_bm25 floors negative idf at EPSILON * average_idf, but on
+            # tiny corpora the average itself can be negative, which would
+            # make the floor negative and silently zero out every match.
+            # Clamp the floor to a small positive value instead.
+            floor = EPSILON * (idf_sum / len(idf))
+            if floor <= 0:
+                floor = EPSILON
+            for term in negatives:
+                idf[term] = floor
+        self.idf = idf
+
+    def scores_for(self, query_code: str) -> Dict[str, float]:
+        """
+        BM25 scores of every symbol against `query_code`.
+
+        Returns only symbols with a positive score. Duplicate query terms
+        contribute once per occurrence (same as BM25Okapi.get_scores).
+        """
+        if not self.ids or self.avgdl == 0:
+            return {}
+        query_tf = Counter(tokenize(query_code))
+        scores: Dict[int, float] = defaultdict(float)
+        for term, q_count in query_tf.items():
+            idf = self.idf.get(term)
+            if idf is None:
+                continue
+            for doc_idx, f in self.postings[term]:
+                denom = f + K1 * (1 - B + B * self.doc_len[doc_idx] / self.avgdl)
+                scores[doc_idx] += q_count * idf * f * (K1 + 1) / denom
+        return {self.ids[i]: s for i, s in scores.items() if s > 0}
+
+
+def get_lexical_index(index) -> LexicalIndex:
+    """
+    Return the RepositoryIndex's lexical index, building and caching it on
+    first use. pipeline.update_index() invalidates the cache when symbols
+    change, so a stale index is never served.
+    """
+    if index._lexical is None:
+        index._lexical = LexicalIndex(index.symbols)
+    return index._lexical

diffcontext/models.py

@@ -0,0 +1,180 @@
+"""
+models.py — Data classes used across the pipeline.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Set
+
+
+@dataclass
+class Symbol:
+    """A single extracted symbol (function/method)."""
+    id: str               # "./path.py:ClassName.method" or "./path.py:func"
+    file: str             # absolute path to source file
+    name: str             # bare name like "ClassName.method" or "func"
+    code: str             # source code text
+    lineno: int = 0       # start line number
+
+
+@dataclass
+class RepositoryIndex:
+    """Complete index of a repository."""
+    symbols: Dict[str, Symbol] = field(default_factory=dict)     # id -> Symbol
+    graph: Dict[str, List[str]] = field(default_factory=dict)    # id -> [dependency ids]
+    broken_files: List[str] = field(default_factory=list)
+
+    # Incremental-update state, populated by pipeline.index_repository().
+    # Private: not part of the public API, excluded from repr/comparison.
+    _repo_path: Optional[str] = field(default=None, repr=False, compare=False)
+    _file_trees: Optional[Dict] = field(default=None, repr=False, compare=False)
+    _import_maps: Optional[Dict] = field(default=None, repr=False, compare=False)
+    _warn_state: Optional[object] = field(default=None, repr=False, compare=False)
+    # Lazily built BM25 index (see lexical.get_lexical_index); invalidated
+    # by update_index() whenever symbols change.
+    _lexical: Optional[object] = field(default=None, repr=False, compare=False)
+    # Lazily built reverse call graph; invalidated by update_index()
+    # whenever the forward graph is rebuilt.
+    _reverse_graph: Optional[Dict[str, Set[str]]] = field(default=None, repr=False, compare=False)
+    # Per-language-adapter edge sets (adapter name -> {id: [dep ids]}),
+    # kept separately so update_index() can rebuild the Python part
+    # without re-running unchanged language adapters. None on a
+    # graph-cache warm start (recomputed on first update that needs it).
+    _lang_graphs: Optional[Dict[str, Dict[str, List[str]]]] = field(default=None, repr=False, compare=False)
+
+    def update(self, changed_files: List[str]) -> "RepositoryIndex":
+        """
+        Incrementally re-index after `changed_files` were edited, created,
+        or deleted. Only those files are re-read and re-parsed; the graph
+        is rebuilt from in-memory ASTs. Mutates and returns this index.
+
+        Only available on indexes created by pipeline.index_repository().
+        """
+        from .pipeline import update_index
+        return update_index(self, changed_files)
+
+    @property
+    def reverse_graph(self) -> Dict[str, Set[str]]:
+        """
+        Reverse graph (callers of each symbol), computed once per index
+        state and cached. Treat the returned dict as read-only: it is
+        shared across callers and only invalidated by update().
+        """
+        if self._reverse_graph is None:
+            rev: Dict[str, Set[str]] = {}
+            for caller, callees in self.graph.items():
+                for callee in callees:
+                    rev.setdefault(callee, set()).add(caller)
+            self._reverse_graph = rev
+        return self._reverse_graph
+
+    @property
+    def total_edges(self) -> int:
+        return sum(len(deps) for deps in self.graph.values())
+
+
+@dataclass
+class DiffResult:
+    """Result of comparing two states."""
+    modified: List[str] = field(default_factory=list)
+    added: List[str] = field(default_factory=list)
+    deleted: List[str] = field(default_factory=list)
+    broken_files: List[str] = field(default_factory=list)
+
+    @property
+    def all_changed(self) -> List[str]:
+        return self.modified + self.added
+
+
+@dataclass
+class ImpactResult:
+    """Result of impact analysis."""
+    changed: List[str] = field(default_factory=list)
+    blast_radius: List[str] = field(default_factory=list)
+    dependencies: List[str] = field(default_factory=list)
+    scores: Dict[str, float] = field(default_factory=dict)
+
+    @property
+    def all_relevant(self) -> List[str]:
+        """All symbols that should be in context, deduplicated, ordered by score."""
+        seen = set()
+        result = []
+        # Score-ordered
+        scored = sorted(self.scores.items(), key=lambda x: x[1], reverse=True)
+        for sym_id, _ in scored:
+            if sym_id not in seen:
+                seen.add(sym_id)
+                result.append(sym_id)
+        # Any remaining that weren't scored
+        for sym_id in self.changed + self.blast_radius + self.dependencies:
+            if sym_id not in seen:
+                seen.add(sym_id)
+                result.append(sym_id)
+        return result
+
+
+@dataclass
+class ContextItem:
+    """
+    One selected symbol, in structured form. The base representation of a
+    compiled context: a harness can filter, reorder, and re-budget these
+    itself instead of consuming the pre-rendered text.
+    """
+    symbol_id: str                 # "./path.py:ClassName.method"
+    code: str                      # full source of the symbol
+    score: float                   # impact score (higher = more relevant)
+    role: str                      # "changed" | "impacted" | "dependency"
+    callers: List[str] = field(default_factory=list)   # full list, untruncated
+    callees: List[str] = field(default_factory=list)   # full list, untruncated
+    token_estimate: int = 0
+
+
+@dataclass
+class ContextPackage:
+    """
+    Final compiled context for an LLM.
+
+    `items` is the structured base representation; `text` is one renderer
+    over it (meta-header + sections + suggestions) for direct LLM pasting.
+    """
+    text: str
+    symbol_count: int
+    token_estimate: int
+    total_repo_tokens: int
+    # Structured selection — the machine-consumable form of `text`'s body.
+    items: List[ContextItem] = field(default_factory=list)
+    # LLM self-awareness fields
+    dropped_symbols: List[str] = field(default_factory=list)   # scored but cut by budget
+    skipped_files: List[str] = field(default_factory=list)     # SyntaxError'd files
+    graph_confidence: float = 1.0                              # fraction of edges that resolved
+
+    @property
+    def reduction_pct(self) -> float:
+        if self.total_repo_tokens == 0:
+            return 0.0
+        return (1 - self.token_estimate / self.total_repo_tokens) * 100
+
+
+@dataclass
+class BenchmarkResult:
+    """Result of a single benchmark run."""
+    name: str
+    repo_path: str
+    changed_functions: List[str]
+    # Graph stats
+    total_symbols: int = 0
+    total_edges: int = 0
+    graph_build_ms: float = 0.0
+    # Retrieval stats
+    retrieved_count: int = 0
+    retrieved_ids: List[str] = field(default_factory=list)
+    # Token stats
+    total_tokens: int = 0
+    context_tokens: int = 0
+    token_reduction_pct: float = 0.0
+    function_reduction_pct: float = 0.0
+    # Precision/Recall (when ground truth available)
+    precision: Optional[float] = None
+    recall: Optional[float] = None
+    f1: Optional[float] = None
+    # Timing
+    pipeline_ms: float = 0.0

diffcontext/parser.py

@@ -0,0 +1,183 @@
+"""
+parser.py — AST-based symbol extraction from Python source files.
+
+Extracts functions, methods (including async), with class-aware naming.
+"""
+
+import ast
+import logging
+import os
+from typing import Dict, List, Optional
+
+from .models import Symbol
+from ._warn_once import warn_syntax_error_once, check_and_warn_encoding
+
+logger = logging.getLogger(__name__)
+
+
+# `def`/`class` are statements, so collection only needs to descend
+# through statement blocks — expression subtrees (the majority of AST
+# nodes) can never contain a definition. Field order mirrors the AST's
+# own field order so collection order matches a full NodeVisitor walk.
+_STMT_BLOCK_FIELDS = ("body", "handlers", "orelse", "finalbody", "cases")
+
+
+def _collect_functions(tree: "ast.Module") -> "List[tuple]":
+    """
+    Collect (qualified_name, node) for every function/method definition,
+    including nested functions, methods of classes defined inside
+    functions, and definitions under conditional blocks (`if
+    TYPE_CHECKING:`, `try/except ImportError`, `match`).
+    """
+    collected: "List[tuple]" = []
+    class_stack: "List[str]" = []
+
+    def _walk(stmts):
+        for node in stmts:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                if class_stack:
+                    name = ".".join(class_stack) + "." + node.name
+                else:
+                    name = node.name
+                collected.append((name, node))
+                _walk(node.body)
+            elif isinstance(node, ast.ClassDef):
+                class_stack.append(node.name)
+                _walk(node.body)
+                class_stack.pop()
+            else:
+                for field in _STMT_BLOCK_FIELDS:
+                    block = getattr(node, field, None)
+                    if block:
+                        _walk(block)
+
+    _walk(tree.body)
+    return collected
+
+
+def _segment_lines(source: str) -> "Optional[List[str]]":
+    """
+    Pre-split source for fast per-symbol segment slicing.
+
+    `ast.get_source_segment` re-splits the ENTIRE file for every symbol —
+    on a large repo that is the single biggest cold-index cost. Splitting
+    once per file and slicing per symbol is equivalent, but only when the
+    file has no `\\r` or `\\f` characters (the parser's line accounting
+    treats those specially); return None then, and the caller falls back
+    to `ast.get_source_segment` for that file.
+    """
+    if "\r" in source or "\f" in source:
+        return None
+    return source.split("\n")
+
+
+def _fast_segment(lines: "List[str]", node) -> "Optional[str]":
+    """Slice a node's source from pre-split lines. AST column offsets are
+    UTF-8 byte offsets, so non-ASCII boundary lines go through bytes."""
+    end_lineno = getattr(node, "end_lineno", None)
+    end_col = getattr(node, "end_col_offset", None)
+    if end_lineno is None or end_col is None:
+        return None
+    lineno = node.lineno - 1
+    end_lineno -= 1
+    col = node.col_offset
+
+    def _cols(line: str, start: "Optional[int]", end: "Optional[int]") -> str:
+        if line.isascii():
+            return line[start:end]
+        return line.encode("utf-8")[start:end].decode("utf-8")
+
+    if end_lineno == lineno:
+        return _cols(lines[lineno], col, end_col)
+    first = _cols(lines[lineno], col, None)
+    last = _cols(lines[end_lineno], None, end_col)
+    return "\n".join([first, *lines[lineno + 1 : end_lineno], last])
+
+
+def extract_symbols(
+    filename: str,
+    repo_path: str,
+    broken_files: "Optional[List[str]]" = None,
+    source: "Optional[str]" = None,
+    tree: "Optional[ast.Module]" = None,
+) -> Dict[str, Symbol]:
+    """
+    Parse a single Python file, return dict of symbol_id -> Symbol.
+
+    Symbol IDs look like: "./relative/path.py:ClassName.method_name"
+
+    If parsing fails and `broken_files` is provided (a list), the file's
+    relative path is appended to it so callers can distinguish "file failed
+    to parse" from "file legitimately has no functions."
+
+    `source` and `tree` may be supplied together to reuse an already-read,
+    already-parsed file (the pipeline parses each file exactly once and
+    shares the result); both must correspond to the same file contents.
+    """
+    relative_file = "./" + os.path.relpath(filename, repo_path)
+
+    if source is None or tree is None:
+        with open(filename, "rb") as f:
+            raw = f.read()
+        check_and_warn_encoding(logger, filename, raw)
+        source = raw.decode("utf-8", errors="ignore")
+
+        try:
+            tree = ast.parse(source)
+        except SyntaxError as e:
+            warn_syntax_error_once(logger, filename, e)
+            if broken_files is not None:
+                broken_files.append(relative_file)
+            return {}
+
+    seg_lines = _segment_lines(source)
+
+    symbols = {}
+    for name, node in _collect_functions(tree):
+        symbol_id = f"{relative_file}:{name}"
+        if seg_lines is not None:
+            try:
+                code = _fast_segment(seg_lines, node)
+            except (IndexError, UnicodeError):
+                code = ast.get_source_segment(source, node)
+        else:
+            code = ast.get_source_segment(source, node)
+        if code is None:
+            continue
+        symbols[symbol_id] = Symbol(
+            id=symbol_id,
+            file=filename,
+            name=name,
+            code=code,
+            lineno=node.lineno,
+        )
+
+    return symbols
+
+
+def extract_all_symbols(
+    repo_path: str,
+    broken_files: "Optional[List[str]]" = None,
+) -> Dict[str, Symbol]:
+    """
+    Extract symbols from all Python files in a repository.
+
+    If `broken_files` is provided (a list), relative paths of any files
+    that failed to parse (SyntaxError) are appended to it.
+    """
+    from .scanner import find_python_files
+    from .cache import SymbolCache
+
+    repo_path = os.path.abspath(repo_path)
+    all_symbols: Dict[str, Symbol] = {}
+    
+    db_path = os.path.join(repo_path, ".diffcontext_cache.db")
+
+    with SymbolCache(db_path) as cache:
+        for filepath in find_python_files(repo_path):
+            def _parse(path: str) -> Dict[str, Symbol]:
+                return extract_symbols(path, repo_path, broken_files=broken_files)
+            
+            all_symbols.update(cache.get_or_parse(filepath, _parse))
+
+    return all_symbols