nabla-cli 0.1.0__py3-none-any.whl

nabla/scanner.py

@@ -0,0 +1,453 @@
+"""T1 — call-site detection + classification. See cli-plan.md T1.
+
+Pure `ast` static analysis. No LLM, no network, no imports of sibling
+nabla modules besides `contract`.
+
+Algorithm, in two passes over every ``*.py`` file under the repo:
+
+1. For every function definition (module-level function or class method,
+   at any nesting depth) find the SDK calls made *directly* in its body
+   (not inside a nested function/class/lambda — those get their own
+   entry when we reach them), and the calls it makes to other
+   project-local functions (resolved through same-file definitions or
+   ``import`` / ``from ... import`` statements, one hop).
+
+2. A function that contains an SDK call directly is a "wrapper
+   candidate". If nothing in the repo calls it, it *is* the site. If
+   one or more project functions call it, each caller becomes its own
+   site (flagged ``wrapper_indirect``) and the wrapper itself is not
+   reported — *unless* the function owns its prompt template (literal
+   or module-constant message content in the call's ``messages``). A
+   template-owning function is the site even when called from elsewhere:
+   its callers pass slot data, not templates.
+"""
+
+from __future__ import annotations
+
+import ast
+import os
+import posixpath
+from dataclasses import dataclass, field
+
+from .contract import (
+    FLAG_ASYNC,
+    FLAG_MULTI_TURN,
+    FLAG_N_GT_1,
+    FLAG_STREAMING,
+    FLAG_VISION,
+    FLAG_WRAPPER_INDIRECT,
+    RawSite,
+    RawSiteRef,
+)
+
+_SKIP_DIR_NAMES = {".git", "node_modules", "__pycache__", ".venv", "venv"}
+
+_VISION_MARKERS = {"image_url", "input_image"}
+
+Span = tuple[int, int]
+FuncKey = tuple[str, str]  # (repo-relative file, qualname)
+
+
+@dataclass
+class _SdkCallInfo:
+    node: ast.Call
+    kind: str
+    flags: list[str]
+    owns_template: bool = False
+
+
+@dataclass
+class _FileImports:
+    # local name -> (target repo-relative file, target function name)
+    names: dict[str, FuncKey] = field(default_factory=dict)
+    # local module alias -> target repo-relative file
+    modules: dict[str, str] = field(default_factory=dict)
+
+
+def scan(repo_path: str) -> list[RawSite]:
+    """Find every OpenAI chat-completions call site under repo_path.
+
+    Resolves within-module client aliases, enumerates wrapper callers
+    (one hop) as distinct sites, detects flags (streaming/async/
+    multi_turn/vision/n_gt_1/wrapper_indirect), and returns a
+    deterministic ordering (sorted by file, then span).
+    """
+    repo_abspath = os.path.abspath(repo_path)
+    files = _discover_python_files(repo_abspath)
+
+    parsed: dict[str, ast.Module] = {}
+    for relpath, abspath in files:
+        with open(abspath, "r", encoding="utf-8") as f:
+            source = f.read()
+        parsed[relpath] = ast.parse(source, filename=abspath)
+
+    top_level_funcs: dict[str, dict[str, ast.AST]] = {
+        relpath: {
+            stmt.name: stmt
+            for stmt in tree.body
+            if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef))
+        }
+        for relpath, tree in parsed.items()
+    }
+
+    imports: dict[str, _FileImports] = {
+        relpath: _resolve_imports(relpath, tree, top_level_funcs)
+        for relpath, tree in parsed.items()
+    }
+
+    direct_sdk: dict[FuncKey, list[_SdkCallInfo]] = {}
+    callee_calls: dict[FuncKey, list[tuple[ast.Call, FuncKey]]] = {}
+
+    for relpath, tree in parsed.items():
+        module_str_consts = {
+            stmt.targets[0].id
+            for stmt in tree.body
+            if isinstance(stmt, ast.Assign)
+            and len(stmt.targets) == 1
+            and isinstance(stmt.targets[0], ast.Name)
+            and isinstance(stmt.value, ast.Constant)
+            and isinstance(stmt.value.value, str)
+        }
+        for func_node, qualname in _iter_all_funcdefs(tree):
+            is_async = isinstance(func_node, ast.AsyncFunctionDef)
+            local_assigns: dict[str, ast.AST] = {}
+            sdk_calls_here: list[_SdkCallInfo] = []
+            callee_here: list[tuple[ast.Call, FuncKey]] = []
+
+            for n in _bounded_walk(func_node):
+                if (
+                    isinstance(n, ast.Assign)
+                    and len(n.targets) == 1
+                    and isinstance(n.targets[0], ast.Name)
+                ):
+                    local_assigns[n.targets[0].id] = n.value
+                if isinstance(n, ast.Call):
+                    kind = _match_sdk_call(n)
+                    if kind is not None:
+                        flags = _call_flags(n, is_async, local_assigns)
+                        sdk_calls_here.append(_SdkCallInfo(
+                            node=n, kind=kind, flags=flags,
+                            owns_template=_owns_template(n, module_str_consts),
+                        ))
+                    else:
+                        target = _resolve_callee(n, relpath, top_level_funcs, imports)
+                        if target is not None:
+                            callee_here.append((n, target))
+
+            key = (relpath, qualname)
+            if sdk_calls_here:
+                direct_sdk[key] = sdk_calls_here
+            if callee_here:
+                callee_calls[key] = callee_here
+
+    wrapper_callers: dict[FuncKey, list[tuple[str, str, ast.Call]]] = {}
+    for (caller_file, caller_func), calls in callee_calls.items():
+        for call_node, target_key in calls:
+            if target_key in direct_sdk:
+                wrapper_callers.setdefault(target_key, []).append(
+                    (caller_file, caller_func, call_node)
+                )
+
+    sites: list[RawSite] = []
+    for key, sdk_calls in direct_sdk.items():
+        file, func = key
+        callers = wrapper_callers.get(key)
+        # A template-owning function is never demoted to a wrapper: its
+        # callers supply slot values, not prompt structure.
+        if callers and all(c.owns_template for c in sdk_calls):
+            callers = None
+        if callers:
+            for sdk_call in sdk_calls:
+                wrapper_ref = RawSiteRef(file=file, span=_span(sdk_call.node), symbol=func)
+                for caller_file, caller_func, call_node in callers:
+                    site_flags = sorted(set(sdk_call.flags) | {FLAG_WRAPPER_INDIRECT})
+                    sites.append(
+                        RawSite(
+                            file=caller_file,
+                            span=_span(call_node),
+                            symbol=caller_func,
+                            kind=sdk_call.kind,
+                            flags=site_flags,
+                            wrapper=wrapper_ref,
+                        )
+                    )
+        else:
+            for sdk_call in sdk_calls:
+                sites.append(
+                    RawSite(
+                        file=file,
+                        span=_span(sdk_call.node),
+                        symbol=func,
+                        kind=sdk_call.kind,
+                        flags=sorted(sdk_call.flags),
+                        wrapper=None,
+                    )
+                )
+
+    sites.sort(key=lambda s: (s.file, s.span))
+    return sites
+
+
+# --------------------------------------------------------------------------
+# Repo discovery
+# --------------------------------------------------------------------------
+
+
+def _discover_python_files(repo_abspath: str) -> list[tuple[str, str]]:
+    results: list[tuple[str, str]] = []
+    for dirpath, dirnames, filenames in os.walk(repo_abspath):
+        if "pyvenv.cfg" in filenames:
+            dirnames[:] = []
+            continue
+        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIR_NAMES]
+        for fname in filenames:
+            if fname.endswith(".py"):
+                abspath = os.path.join(dirpath, fname)
+                relpath = os.path.relpath(abspath, repo_abspath).replace(os.sep, "/")
+                results.append((relpath, abspath))
+    results.sort()
+    return results
+
+
+# --------------------------------------------------------------------------
+# Scope-bounded traversal
+# --------------------------------------------------------------------------
+
+
+def _bounded_walk(func_node: ast.AST):
+    """Yield every descendant in func_node's body without crossing into a
+    nested function/class/lambda scope (that scope is walked separately
+    when we reach it as its own top-level entry from `_iter_all_funcdefs`)."""
+
+    def walk(n: ast.AST):
+        yield n
+        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+            return
+        for child in ast.iter_child_nodes(n):
+            yield from walk(child)
+
+    for stmt in func_node.body:
+        yield from walk(stmt)
+
+
+def _iter_all_funcdefs(tree: ast.Module):
+    """Yield (func_node, qualname) for every function/method defined
+    anywhere in the module. qualname is "Class.method" for methods
+    (one level), else the bare function name."""
+
+    def walk(node: ast.AST, class_stack: list[str]):
+        for child in ast.iter_child_nodes(node):
+            if isinstance(child, ast.ClassDef):
+                yield from walk(child, class_stack + [child.name])
+            elif isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                qualname = ".".join(class_stack + [child.name])
+                yield (child, qualname)
+                yield from walk(child, class_stack)
+            else:
+                yield from walk(child, class_stack)
+
+    yield from walk(tree, [])
+
+
+# --------------------------------------------------------------------------
+# SDK call matching + flag detection
+# --------------------------------------------------------------------------
+
+
+def _match_sdk_call(call: ast.Call) -> str | None:
+    """Match `<base...>.chat.completions.create(...)` -> "create" or
+    `<base...>.beta.chat.completions.parse(...)` -> "parse", regardless
+    of what `<base...>` resolves to (client alias name, `self.client`,
+    etc.)."""
+    func = call.func
+    if not isinstance(func, ast.Attribute):
+        return None
+
+    if func.attr == "create":
+        completions = func.value
+        if isinstance(completions, ast.Attribute) and completions.attr == "completions":
+            chat = completions.value
+            if isinstance(chat, ast.Attribute) and chat.attr == "chat":
+                return "create"
+        return None
+
+    if func.attr == "parse":
+        completions = func.value
+        if isinstance(completions, ast.Attribute) and completions.attr == "completions":
+            chat = completions.value
+            if isinstance(chat, ast.Attribute) and chat.attr == "chat":
+                beta = chat.value
+                if isinstance(beta, ast.Attribute) and beta.attr == "beta":
+                    return "parse"
+        return None
+
+    return None
+
+
+def _get_kwarg(call: ast.Call, name: str) -> ast.AST | None:
+    for kw in call.keywords:
+        if kw.arg == name:
+            return kw.value
+    return None
+
+
+def _call_flags(call: ast.Call, is_async: bool, local_assigns: dict[str, ast.AST]) -> list[str]:
+    flags: list[str] = []
+    if is_async:
+        flags.append(FLAG_ASYNC)
+
+    stream_val = _get_kwarg(call, "stream")
+    if isinstance(stream_val, ast.Constant) and stream_val.value is True:
+        flags.append(FLAG_STREAMING)
+
+    n_val = _get_kwarg(call, "n")
+    if (
+        isinstance(n_val, ast.Constant)
+        and isinstance(n_val.value, int)
+        and not isinstance(n_val.value, bool)
+        and n_val.value >= 2
+    ):
+        flags.append(FLAG_N_GT_1)
+
+    messages_val = _get_kwarg(call, "messages")
+    if messages_val is not None:
+        if _is_multi_turn(messages_val, local_assigns):
+            flags.append(FLAG_MULTI_TURN)
+        if _has_vision_content(messages_val):
+            flags.append(FLAG_VISION)
+
+    return flags
+
+
+def _owns_template(call: ast.Call, module_str_consts: set[str]) -> bool:
+    """True when the SDK call's ``messages`` literal carries template text of
+    its own — a string-literal content value, a module-level string constant,
+    or an f-string with non-trivial literal fragments. Distinguishes a
+    template-owning site (e.g. a fixed system prompt + one data parameter)
+    from a generic wrapper whose entire prompt arrives via parameters."""
+    messages = _get_kwarg(call, "messages")
+    if not isinstance(messages, ast.List):
+        return False
+    for elt in messages.elts:
+        if not isinstance(elt, ast.Dict):
+            continue
+        for k, v in zip(elt.keys, elt.values):
+            if not (isinstance(k, ast.Constant) and k.value == "content"):
+                continue
+            if isinstance(v, ast.Constant) and isinstance(v.value, str) and v.value.strip():
+                return True
+            if isinstance(v, ast.Name) and v.id in module_str_consts:
+                return True
+            if isinstance(v, ast.JoinedStr) and any(
+                isinstance(part, ast.Constant)
+                and isinstance(part.value, str)
+                and part.value.strip()
+                for part in v.values
+            ):
+                return True
+    return False
+
+
+def _is_multi_turn(value_node: ast.AST, local_assigns: dict[str, ast.AST]) -> bool:
+    """True when `messages` is derived from a parameter or accumulated
+    history rather than being a list literal built in-function."""
+    if isinstance(value_node, ast.List):
+        return False
+    if isinstance(value_node, ast.Name):
+        assigned = local_assigns.get(value_node.id)
+        if assigned is None:
+            # Not locally assigned -> comes from an outer scope (a
+            # parameter, module global, etc.).
+            return True
+        return not isinstance(assigned, ast.List)
+    return True
+
+
+def _has_vision_content(node: ast.AST) -> bool:
+    for n in ast.walk(node):
+        if isinstance(n, ast.Constant) and isinstance(n.value, str) and n.value in _VISION_MARKERS:
+            return True
+    return False
+
+
+def _span(node: ast.AST) -> Span:
+    return (node.lineno, node.end_lineno)
+
+
+# --------------------------------------------------------------------------
+# Import resolution (same-repo, one hop)
+# --------------------------------------------------------------------------
+
+
+def _resolve_imports(
+    relpath: str,
+    tree: ast.Module,
+    top_level_funcs: dict[str, dict[str, ast.AST]],
+) -> _FileImports:
+    result = _FileImports()
+    file_dir = posixpath.dirname(relpath)
+
+    for stmt in tree.body:
+        if isinstance(stmt, ast.ImportFrom):
+            module_parts = stmt.module.split(".") if stmt.module else []
+            if stmt.level == 0:
+                candidate = "/".join(module_parts) + ".py"
+            else:
+                target_dir = file_dir
+                for _ in range(stmt.level - 1):
+                    target_dir = posixpath.dirname(target_dir)
+                parts = [p for p in [target_dir, *module_parts] if p]
+                candidate = "/".join(parts) + ".py"
+
+            if candidate in top_level_funcs:
+                for alias in stmt.names:
+                    orig_name = alias.name
+                    local_name = alias.asname or alias.name
+                    if orig_name in top_level_funcs[candidate]:
+                        result.names[local_name] = (candidate, orig_name)
+                    else:
+                        # Might be a submodule/name we can't resolve
+                        # (e.g. `from . import module`); ignore.
+                        pass
+            # Track the module itself too, for `from .llm_utils import
+            # x` used elsewhere as attribute access — not needed here.
+
+        elif isinstance(stmt, ast.Import):
+            for alias in stmt.names:
+                module_parts = alias.name.split(".")
+                candidate = "/".join(module_parts) + ".py"
+                if candidate in top_level_funcs:
+                    local_name = alias.asname or module_parts[0]
+                    result.modules[local_name] = candidate
+
+    return result
+
+
+def _resolve_callee(
+    call: ast.Call,
+    relpath: str,
+    top_level_funcs: dict[str, dict[str, ast.AST]],
+    imports: dict[str, _FileImports],
+) -> FuncKey | None:
+    """Resolve a call to another project-local function, one hop, via
+    same-file definitions or this file's imports. Returns None for
+    anything else (external libraries, dynamic dispatch, etc.)."""
+    func = call.func
+    file_imports = imports.get(relpath, _FileImports())
+
+    if isinstance(func, ast.Name):
+        name = func.id
+        if name in top_level_funcs.get(relpath, {}):
+            return (relpath, name)
+        if name in file_imports.names:
+            return file_imports.names[name]
+        return None
+
+    if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):
+        base = func.value.id
+        target_file = file_imports.modules.get(base)
+        if target_file is not None and func.attr in top_level_funcs.get(target_file, {}):
+            return (target_file, func.attr)
+        return None
+
+    return None