coltrace 0.1.0.dev4__py3-none-any.whl

coltrace/collector.py

@@ -0,0 +1,561 @@
+"""Collect syntactic occurrences and conservative evidence about their receivers.
+
+This is not a lineage engine: no expression dependencies, column versions or
+interprocedural effects are inferred here. Unsupported contexts stay explicit.
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+from collections import Counter
+from dataclasses import dataclass
+from typing import Literal
+
+from .models import (
+    ColumnOccurrence,
+    DataFrameInfo,
+    Diagnostic,
+    Evidence,
+    FunctionInfo,
+    Location,
+)
+
+
+@dataclass(frozen=True)
+class _Binding:
+    evidence: Evidence = "unresolved"
+    reason: str = "No supported DataFrame origin or annotation was resolved."
+    qualified: str | None = None
+    dataframe_id: str | None = None
+
+
+class _BoundNames(ast.NodeVisitor):
+    """Find bindings in one Python scope without entering nested scopes."""
+
+    def __init__(self) -> None:
+        self.names: Counter[str] = Counter()
+
+    def visit_Name(self, node: ast.Name) -> None:
+        if isinstance(node.ctx, (ast.Store, ast.Del)):
+            self.names[node.id] += 1
+
+    def visit_Import(self, node: ast.Import) -> None:
+        for item in node.names:
+            self.names[item.asname or item.name.split(".")[0]] += 1
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        for item in node.names:
+            self.names[item.asname or item.name] += 1
+
+    def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
+        self.names[node.name] += 1
+
+    visit_AsyncFunctionDef = visit_FunctionDef
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        self.names[node.name] += 1
+
+    def visit_Lambda(self, node: ast.Lambda) -> None:
+        pass
+
+    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None:
+        if node.name:
+            self.names[node.name] += 1
+        self.generic_visit(node)
+
+    def visit_MatchAs(self, node: ast.MatchAs) -> None:
+        if node.name:
+            self.names[node.name] += 1
+        self.generic_visit(node)
+
+    def visit_MatchStar(self, node: ast.MatchStar) -> None:
+        if node.name:
+            self.names[node.name] += 1
+
+    def visit_MatchMapping(self, node: ast.MatchMapping) -> None:
+        if node.rest:
+            self.names[node.rest] += 1
+        self.generic_visit(node)
+
+    def visit_ListComp(self, node: ast.AST) -> None:
+        pass
+
+    visit_SetComp = visit_ListComp
+    visit_DictComp = visit_ListComp
+    visit_GeneratorExp = visit_ListComp
+
+
+def _parameters(args: ast.arguments) -> list[ast.arg]:
+    result = args.posonlyargs + args.args + args.kwonlyargs
+    if args.vararg:
+        result.append(args.vararg)
+    if args.kwarg:
+        result.append(args.kwarg)
+    return result
+
+
+def _literal(node: ast.AST) -> bool:
+    if isinstance(node, ast.Constant):
+        return True
+    if isinstance(node, (ast.List, ast.Tuple)):
+        return all(_literal(item) for item in node.elts)
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
+        return _literal(node.operand)
+    return False
+
+
+class OccurrenceCollector(ast.NodeVisitor):
+    def __init__(self, path: str, module: str, source: str) -> None:
+        self.path = path
+        self.module = module
+        self.source = source
+        self.lines = source.split("\n")
+        self.scope = "<module>"
+        self.env: dict[str, _Binding] = {}
+        self._module_env = self.env
+        self._module_bindings: Counter[str] = Counter()
+        self._uncertain = 0
+        self.occurrences: list[ColumnOccurrence] = []
+        self.functions: list[FunctionInfo] = []
+        self.dataframes: list[DataFrameInfo] = []
+        self.diagnostics: list[Diagnostic] = []
+
+    def _location(self, node: ast.AST) -> Location:
+        def column(line: int, byte_offset: int) -> int:
+            return len(self.lines[line - 1].encode("utf-8")[:byte_offset].decode("utf-8"))
+
+        line = node.lineno  # type: ignore[attr-defined]
+        end_line = getattr(node, "end_lineno", line) or line
+        return Location(
+            self.path,
+            line,
+            column(line, node.col_offset),  # type: ignore[attr-defined]
+            end_line,
+            column(end_line, getattr(node, "end_col_offset", 0) or 0),
+        )
+
+    def _diagnose(self, code: str, node: ast.AST, message: str, hint: str) -> None:
+        self.diagnostics.append(
+            Diagnostic(
+                code,
+                "warning",
+                message,
+                self._location(node),
+                self.scope,
+                hint,
+            )
+        )
+
+    def _qualified(self, node: ast.AST, env: dict[str, _Binding] | None = None) -> str | None:
+        bindings = self.env if env is None else env
+        if isinstance(node, ast.Name):
+            return bindings.get(node.id, _Binding()).qualified
+        if isinstance(node, ast.Attribute):
+            base = self._qualified(node.value, bindings)
+            return f"{base}.{node.attr}" if base else None
+        return None
+
+    def _dataframe(self, node: ast.AST, origin: str, evidence: Evidence) -> _Binding:
+        if self._uncertain:
+            return _Binding(reason="The receiver is in an unsupported execution context.")
+        location = self._location(node)
+        key = f"{self.path}:{self.scope}:{location.line}:{location.column}:{origin}"
+        identifier = "df:" + hashlib.sha256(key.encode()).hexdigest()[:24]
+        reason = (
+            f"Resolved Pandas origin: {origin}."
+            if evidence == "confirmed"
+            else "DataFrame annotation is an unverified type assumption."
+        )
+        self.dataframes.append(
+            DataFrameInfo(
+                identifier,
+                self.scope,
+                origin,
+                evidence,
+                reason,
+                location,
+            )
+        )
+        return _Binding(evidence, reason, dataframe_id=identifier)
+
+    def _infer(self, node: ast.AST) -> _Binding:
+        if isinstance(node, ast.Name):
+            return self.env.get(node.id, _Binding())
+        if isinstance(node, (ast.Constant, ast.Dict, ast.List, ast.Tuple, ast.Set)):
+            return _Binding("rejected", "The receiver is a literal, not a DataFrame.")
+        if isinstance(node, ast.Call):
+            origin = self._qualified(node.func)
+            if origin in {"pandas.read_csv", "pandas.read_parquet"}:
+                if origin == "pandas.read_csv" and (
+                    len(node.args) > 1
+                    or any(isinstance(arg, ast.Starred) for arg in node.args)
+                    or any(
+                        kw.arg is None
+                        or (
+                            kw.arg == "iterator"
+                            and not (isinstance(kw.value, ast.Constant) and kw.value.value is False)
+                        )
+                        or (
+                            kw.arg == "chunksize"
+                            and not (isinstance(kw.value, ast.Constant) and kw.value.value is None)
+                        )
+                        for kw in node.keywords
+                    )
+                ):
+                    return _Binding(
+                        reason="CSV reader options may return an iterator, not a DataFrame."
+                    )
+                return self._dataframe(node, origin, "confirmed")
+            if origin == "pandas.DataFrame":
+                data = (
+                    node.args[0]
+                    if len(node.args) == 1
+                    else next(
+                        (kw.value for kw in node.keywords if kw.arg == "data"),
+                        None,
+                    )
+                )
+                if (
+                    isinstance(data, ast.Dict)
+                    and all(
+                        isinstance(key, ast.Constant) and isinstance(key.value, str)
+                        for key in data.keys
+                    )
+                    and all(_literal(value) for value in data.values)
+                ):
+                    return self._dataframe(node, origin, "confirmed")
+        return _Binding()
+
+    def _annotation(
+        self,
+        annotation: ast.expr | None,
+        node: ast.AST,
+        env: dict[str, _Binding],
+    ) -> _Binding:
+        # Parsing a quoted annotation only inspects syntax; no eval or import.
+        if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str):
+            try:
+                annotation = ast.parse(annotation.value, mode="eval").body
+            except SyntaxError:
+                return _Binding()
+        if annotation and self._qualified(annotation, env) == "pandas.DataFrame":
+            return self._dataframe(node, "annotation:pandas.DataFrame", "inferred")
+        return _Binding()
+
+    def visit_Module(self, node: ast.Module) -> None:
+        bindings = _BoundNames()
+        bindings.visit(node)
+        self._module_bindings = bindings.names
+        self.generic_visit(node)
+
+    def visit_Import(self, node: ast.Import) -> None:
+        for item in node.names:
+            name = item.asname or item.name.split(".")[0]
+            self.env[name] = _Binding(qualified="pandas") if item.name == "pandas" else _Binding()
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        for item in node.names:
+            name = item.asname or item.name
+            known = (
+                node.module == "pandas"
+                and node.level == 0
+                and item.name
+                in {
+                    "DataFrame",
+                    "read_csv",
+                    "read_parquet",
+                }
+            )
+            self.env[name] = _Binding(qualified=f"pandas.{item.name}") if known else _Binding()
+        if any(item.name == "*" for item in node.names):
+            self._diagnose(
+                "CT_STAR_IMPORT",
+                node,
+                "Star imports cannot be resolved in Phase 1.",
+                "Use explicit imports.",
+            )
+            for name in self.env:
+                self.env[name] = _Binding(reason="A star import can shadow existing bindings.")
+
+    def _bind_target(self, target: ast.AST, binding: _Binding) -> None:
+        if isinstance(target, ast.Name):
+            self.env[target.id] = binding
+        elif isinstance(target, (ast.Tuple, ast.List)):
+            for item in target.elts:
+                self._bind_target(item, _Binding(reason="Unpacking is not resolved in Phase 1."))
+        elif isinstance(target, ast.Attribute):
+            root = target.value
+            while isinstance(root, ast.Attribute):
+                root = root.value
+            if isinstance(root, ast.Name):
+                affected = self.env.get(root.id, _Binding())
+                if affected.qualified:
+                    for name, value in list(self.env.items()):
+                        if value.qualified and value.qualified.startswith("pandas"):
+                            self.env[name] = _Binding(reason="A Pandas API attribute was modified.")
+                self._invalidate({affected.dataframe_id})
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        self.visit(node.value)
+        binding = self._infer(node.value)
+        for target in node.targets:
+            self.visit(target)
+            self._bind_target(target, binding)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        if node.value:
+            self.visit(node.value)
+            binding = self._infer(node.value)
+        else:
+            binding = (
+                self.env.get(node.target.id, _Binding())
+                if isinstance(node.target, ast.Name)
+                else _Binding()
+            )
+        if binding.evidence == "unresolved":
+            binding = self._annotation(node.annotation, node, self.env)
+        self.visit(node.target)
+        self._bind_target(node.target, binding)
+
+    def visit_Delete(self, node: ast.Delete) -> None:
+        for target in node.targets:
+            self.visit(target)
+            self._bind_target(target, _Binding(reason="This binding was deleted."))
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> None:
+        if isinstance(node.target, ast.Subscript):
+            self._subscript(node.target, "read_write", force_unsupported=True)
+        else:
+            self.visit(node.target)
+        self.visit(node.value)
+        self._bind_target(node.target, _Binding(reason="Augmented assignments are unsupported."))
+
+    def visit_FunctionDef(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
+        parent_env, parent_scope, uncertain = self.env, self.scope, self._uncertain
+        for expression in [*node.decorator_list, *node.args.defaults]:
+            self.visit(expression)
+        for default in node.args.kw_defaults:
+            if default:
+                self.visit(default)
+        scope = node.name if parent_scope == "<module>" else f"{parent_scope}.{node.name}"
+        parameters = _parameters(node.args)
+        self.functions.append(
+            FunctionInfo(
+                node.name,
+                scope,
+                self.module,
+                tuple(arg.arg for arg in parameters),
+                isinstance(node, ast.AsyncFunctionDef),
+                self._location(node),
+            )
+        )
+        locals_ = _BoundNames()
+        for statement in node.body:
+            locals_.visit(statement)
+        self.env = {
+            name: binding
+            for name, binding in self._module_env.items()
+            if binding.qualified and self._module_bindings[name] == 1
+        }
+        self.env.update({name: _Binding() for name in locals_.names})
+        self.scope = scope
+        self._uncertain += int(
+            parent_scope != "<module>"
+            or bool(node.decorator_list)
+            or isinstance(node, ast.AsyncFunctionDef)
+        )
+        for arg in parameters:
+            if arg is node.args.vararg or arg is node.args.kwarg:
+                self.env[arg.arg] = _Binding(reason="Variadic parameters are containers.")
+            else:
+                self.env[arg.arg] = self._annotation(arg.annotation, arg, parent_env)
+        if self._uncertain:
+            self._diagnose(
+                "CT_FUNCTION_CONTEXT",
+                node,
+                "Nested, decorated, async or conditional functions have unresolved context.",
+                "Occurrences are candidates; direct function analysis is a later checkpoint.",
+            )
+        for statement in node.body:
+            self.visit(statement)
+        self.env, self.scope, self._uncertain = parent_env, parent_scope, uncertain
+        self.env[node.name] = _Binding()
+
+    visit_AsyncFunctionDef = visit_FunctionDef
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        env, scope, uncertain = self.env, self.scope, self._uncertain
+        for expression in [*node.decorator_list, *node.bases]:
+            self.visit(expression)
+        for keyword in node.keywords:
+            self.visit(keyword.value)
+        self.env = dict(env)
+        self.scope = node.name if scope == "<module>" else f"{scope}.{node.name}"
+        self._uncertain += 1
+        for statement in node.body:
+            self.visit(statement)
+        self.env, self.scope, self._uncertain = env, scope, uncertain
+        self.env[node.name] = _Binding()
+
+    def visit_Lambda(self, node: ast.Lambda) -> None:
+        env, scope, uncertain = self.env, self.scope, self._uncertain
+        for default in [*node.args.defaults, *node.args.kw_defaults]:
+            if default:
+                self.visit(default)
+        self.env = dict(env)
+        self.env.update({arg.arg: _Binding() for arg in _parameters(node.args)})
+        self.scope = f"{scope}.<lambda>@{node.lineno}:{node.col_offset}"
+        self._uncertain += 1
+        self.visit(node.body)
+        self.env, self.scope, self._uncertain = env, scope, uncertain
+
+    def _invalidate(self, ids: set[str | None]) -> None:
+        for name, binding in list(self.env.items()):
+            if binding.dataframe_id and binding.dataframe_id in ids:
+                self.env[name] = _Binding(reason="A potentially mutating operation is unresolved.")
+
+    def visit_Call(self, node: ast.Call) -> None:
+        self.generic_visit(node)
+        origin = self._qualified(node.func)
+        if origin in {"pandas.read_csv", "pandas.read_parquet", "pandas.DataFrame"}:
+            return
+        inputs: list[ast.AST] = [*node.args, *(kw.value for kw in node.keywords)]
+        if isinstance(node.func, ast.Attribute):
+            inputs.append(node.func.value)
+        ids = {
+            self.env.get(arg.id, _Binding()).dataframe_id
+            for arg in inputs
+            if isinstance(arg, ast.Name)
+        } - {None}
+        if ids:
+            self._diagnose(
+                "CT_OPAQUE_CALL",
+                node,
+                "Effects or returned DataFrame of this call are not resolved in Phase 1.",
+                "Receiver evidence after this call is invalidated, including direct aliases.",
+            )
+            self._invalidate(ids)
+
+    def _subscript(
+        self,
+        node: ast.Subscript,
+        access: Literal["read", "write", "delete", "read_write"] | None = None,
+        *,
+        force_unsupported: bool = False,
+    ) -> None:
+        receiver, label = node.value, node.slice
+        syntax: Literal["subscript", "loc"] = "subscript"
+        full_rows = True
+        if isinstance(receiver, ast.Attribute) and receiver.attr == "loc":
+            receiver = receiver.value
+            syntax = "loc"
+            full_rows = False
+            if isinstance(label, ast.Tuple) and len(label.elts) == 2:
+                rows, label = label.elts
+                full_rows = isinstance(rows, ast.Slice) and all(
+                    value is None for value in (rows.lower, rows.upper, rows.step)
+                )
+        binding = (
+            self.env.get(receiver.id, _Binding()) if isinstance(receiver, ast.Name) else _Binding()
+        )
+        if self._uncertain and binding.evidence != "rejected":
+            binding = _Binding(reason="This occurrence is in an unsupported execution context.")
+        column = (
+            label.value
+            if isinstance(label, ast.Constant) and isinstance(label.value, str)
+            else None
+        )
+        supported = column is not None and full_rows and not force_unsupported
+        mode: Literal["read", "write", "delete", "read_write"] = access or (
+            "write"
+            if isinstance(node.ctx, ast.Store)
+            else "delete"
+            if isinstance(node.ctx, ast.Del)
+            else "read"
+        )
+        self.occurrences.append(
+            ColumnOccurrence(
+                column,
+                ast.get_source_segment(self.source, receiver) or "<expression>",
+                binding.dataframe_id,
+                binding.evidence,
+                binding.reason,
+                mode,
+                syntax,
+                supported,
+                self.scope,
+                self.module,
+                self._location(node),
+            )
+        )
+        if binding.evidence == "unresolved":
+            self._diagnose(
+                "CT_UNKNOWN_RECEIVER",
+                node,
+                "The subscript receiver is not a resolved DataFrame.",
+                "Provide a recognized Pandas origin or an explicit DataFrame annotation.",
+            )
+        elif binding.evidence != "rejected" and not supported:
+            self._diagnose(
+                "CT_UNSUPPORTED_INDEX",
+                node,
+                "Only literal string columns and full-row .loc accesses are supported.",
+                "Dynamic labels, partial rows and augmented writes remain unresolved.",
+            )
+        self.generic_visit(node)
+
+    def visit_Subscript(self, node: ast.Subscript) -> None:
+        self._subscript(node)
+
+    def _unsupported_context(self, node: ast.AST) -> None:
+        env, uncertain = self.env, self._uncertain
+        self.env = dict(env)
+        self._uncertain += 1
+        self._diagnose(
+            "CT_CONTROL_FLOW",
+            node,
+            "Control flow is not resolved in Phase 1.",
+            "Occurrences in this region are candidates; affected evidence is invalidated.",
+        )
+        self.generic_visit(node)
+        self.env, self._uncertain = env, uncertain
+        names = _BoundNames()
+        names.visit(node)
+        for name in names.names:
+            self.env[name] = _Binding(reason="The binding is affected by unsupported control flow.")
+        self._invalidate({binding.dataframe_id for binding in self.env.values()})
+        for name, binding in list(self.env.items()):
+            if binding.qualified:
+                self.env[name] = _Binding(
+                    reason="API bindings across this boundary are unresolved."
+                )
+
+    visit_If = _unsupported_context
+    visit_For = _unsupported_context
+    visit_AsyncFor = _unsupported_context
+    visit_While = _unsupported_context
+    visit_Try = _unsupported_context
+    visit_TryStar = _unsupported_context
+    visit_With = _unsupported_context
+    visit_AsyncWith = _unsupported_context
+    visit_Match = _unsupported_context
+    visit_ListComp = _unsupported_context
+    visit_SetComp = _unsupported_context
+    visit_DictComp = _unsupported_context
+    visit_GeneratorExp = _unsupported_context
+    visit_NamedExpr = _unsupported_context
+
+    def visit_Global(self, node: ast.Global | ast.Nonlocal) -> None:
+        for name in node.names:
+            self.env[name] = _Binding(
+                reason="Global/nonlocal bindings are not resolved in Phase 1."
+            )
+        self._diagnose(
+            "CT_EXTERNAL_BINDING",
+            node,
+            "Global/nonlocal state is not resolved.",
+            "Prefer explicit parameters for statically analyzable inputs.",
+        )
+
+    visit_Nonlocal = visit_Global