stackthread 0.1.0.dev2__py3-none-any.whl

stackthread/lang/python/resolve.py

@@ -0,0 +1,448 @@
+"""Name resolution across a Python source (design §8.3).
+
+Works only from facts, so it is fast and testable. Resolution follows imports and re-exports through packages,
+`self`/`cls`/`super()`, class bases inside the project, and simple local type hints (`x = Foo()`, `db: Session`,
+`with Session() as s`, `x = Order.objects.get(...)`).
+"""
+
+from __future__ import annotations
+
+import re
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Literal
+
+RefKind = Literal["def", "module", "var", "external"]
+
+#: Callee names too common to resolve by name alone; they stay unresolved unless the receiver's type is known.
+COMMON_METHODS = frozenset(
+    {
+        "get",
+        "set",
+        "update",
+        "save",
+        "delete",
+        "add",
+        "append",
+        "extend",
+        "pop",
+        "remove",
+        "items",
+        "keys",
+        "values",
+        "copy",
+        "format",
+        "join",
+        "split",
+        "strip",
+        "replace",
+        "filter",
+        "all",
+        "first",
+        "last",
+        "count",
+        "create",
+        "exists",
+        "order_by",
+        "select_related",
+        "prefetch_related",
+        "annotate",
+        "aggregate",
+        "values_list",
+        "exclude",
+        "distinct",
+        "only",
+        "defer",
+        "read",
+        "write",
+        "close",
+        "open",
+        "send",
+        "run",
+        "start",
+        "stop",
+        "call",
+        "apply",
+        "execute",
+        "commit",
+        "flush",
+        "query",
+        "load",
+        "dump",
+        "dumps",
+        "loads",
+        "encode",
+        "decode",
+        "lower",
+        "upper",
+        "startswith",
+        "endswith",
+        "setdefault",
+        "sort",
+        "reverse",
+        "clear",
+        "insert",
+        "index",
+        "find",
+        "match",
+        "search",
+        "sub",
+        "group",
+        "validate",
+        "clean",
+        "is_valid",
+        "render",
+        "dispatch",
+        "handle",
+        "process",
+        "parse",
+        "build",
+        "check",
+        "init",
+        "setup",
+        "reset",
+        "refresh",
+        "next",
+        "list",
+        "retrieve",
+        "destroy",
+        "partial_update",
+        "post",
+        "put",
+        "patch",
+        "head",
+        "options",
+        "info",
+        "debug",
+        "warning",
+        "error",
+        "exception",
+        "critical",
+        "log",
+        "emit",
+        "wait",
+        "result",
+        "submit",
+        "map",
+        "__init__",
+        "__call__",
+        "__enter__",
+        "__exit__",
+        "__str__",
+        "__repr__",
+    }
+)
+
+_SEGMENT = re.compile(r"[^.\[\]()]+(?:\(\))?(?:\[[^\]]*\])*")
+
+
+@dataclass(frozen=True, slots=True)
+class Ref:
+    kind: RefKind
+    key: str  # def: "module:qualname"; module: dotted name; var: "module:name"; external: dotted path
+
+    @property
+    def module(self) -> str:
+        return self.key.split(":", 1)[0] if self.kind in ("def", "var") else self.key
+
+
+def split_dotted(path: str) -> list[str]:
+    """Split a callee path into segments.
+
+    `Order.objects.filter().update` -> [Order, objects, filter(), update]; `db[orders].find` -> [db[orders], find].
+    """
+    return _SEGMENT.findall(path)
+
+
+def bare(segment: str) -> str:
+    return segment.split("(", 1)[0].split("[", 1)[0]
+
+
+class PyIndex:
+    """Every module of one Python source, with symbol tables and lookups."""
+
+    def __init__(self, facts: list[dict[str, Any]]) -> None:
+        self.modules: dict[str, dict[str, Any]] = {}
+        self.defs: dict[str, dict[str, Any]] = {}
+        self.symbols: dict[str, dict[str, tuple[str, Any]]] = defaultdict(dict)
+        self.stars: dict[str, list[str]] = defaultdict(list)
+        self.vars: dict[str, dict[str, Any]] = {}
+        self.methods_by_name: dict[str, list[str]] = defaultdict(list)
+        self.children: dict[str, dict[str, str]] = defaultdict(dict)
+        self._resolving: set[str] = set()
+        for f in facts:
+            if f["module"]:
+                self.modules[f["module"]] = f
+        for module, f in self.modules.items():
+            table = self.symbols[module]
+            for imp in f["imports"]:
+                if imp["name"] == "*":
+                    self.stars[module].append(imp["module"])
+                elif not imp.get("lazy") or imp["local"] not in table:
+                    table[imp["local"]] = ("import", imp)
+            for a in f["assigns"]:
+                key = f"{module}:{a['name']}"
+                if a["aug"] and key in self.vars:
+                    self.vars[key].setdefault("augments", []).append(a)
+                    continue
+                self.vars[key] = dict(a, module=module)
+                table[a["name"]] = ("var", key)
+            for d in f["defs"]:
+                key = f"{module}:{d['q']}"
+                d = dict(d, module=module, file=f["file"], key=key)
+                self.defs[key] = d
+                if d["parent"] is None:
+                    table[d["name"]] = ("def", key)
+                else:
+                    self.children[f"{module}:{d['parent']}"][d["name"]] = key
+                if d["kind"] == "method" and not d["name"].startswith("__"):
+                    self.methods_by_name[d["name"]].append(key)
+
+    # ---- names ----
+
+    def resolve_name(self, module: str, name: str, _seen: frozenset[tuple[str, str]] = frozenset()) -> Ref | None:
+        if (module, name) in _seen:
+            return None
+        seen = _seen | {(module, name)}
+        entry = self.symbols.get(module, {}).get(name)
+        if entry is None:
+            if f"{module}.{name}" in self.modules:  # a submodule reached as a package attribute
+                return Ref("module", f"{module}.{name}")
+            for star in self.stars.get(module, []):
+                if star in self.modules and (found := self.resolve_name(star, name, seen)):
+                    return found
+            return None
+        kind, data = entry
+        if kind == "def":
+            return Ref("def", data)
+        if kind == "var":
+            return Ref("var", data)
+        imp = data
+        target, imported = imp["module"], imp["name"]
+        if imported is None:
+            return Ref("module", target) if target in self.modules else Ref("external", target)
+        if f"{target}.{imported}" in self.modules:
+            return Ref("module", f"{target}.{imported}")
+        if target in self.modules:
+            return self.resolve_name(target, imported, seen) or Ref("external", f"{target}.{imported}")
+        return Ref("external", f"{target}.{imported}" if target else imported)
+
+    def resolve_dotted(self, module: str, path: str, scope: dict[str, Any] | None = None) -> Ref | None:
+        """Resolve `a.b.c` from inside `module` (and optionally a function scope with local type hints)."""
+        segments = split_dotted(path)
+        if not segments:
+            return None
+        head = bare(segments[0])
+        ref: Ref | None = None
+        if scope is not None:
+            ref = self._scope_head(module, head, scope)
+        if ref is None:
+            ref = self.resolve_name(module, head)
+        if ref is None:
+            return None
+        for seg in segments[1:]:
+            ref = self.member(ref, bare(seg))
+            if ref is None:
+                return None
+        return ref
+
+    def _scope_head(self, module: str, head: str, scope: dict[str, Any]) -> Ref | None:
+        if head in ("self", "cls"):
+            cls = self.enclosing_class(scope)
+            return Ref("def", cls["key"]) if cls else None
+        hint = self.local_hint(scope, head)
+        if hint is None:
+            return None
+        return self.instance_type(module, hint, scope)
+
+    def local_hint(self, scope: dict[str, Any], name: str) -> str | None:
+        current: dict[str, Any] | None = scope
+        while current is not None:
+            if name in current.get("locals", {}):
+                return str(current["locals"][name])
+            parent = current.get("parent")
+            current = self.defs.get(f"{current['module']}:{parent}") if parent else None
+        return None
+
+    def instance_type(self, module: str, hint: str, scope: dict[str, Any] | None = None) -> Ref | None:
+        """What an instance described by a hint is: `Foo()` -> class Foo; `Session` -> class Session;
+        `make_service()` -> its return annotation; `Order.objects.get()` -> model Order."""
+        orm = re.match(r"^([\w.]+)\.(?:objects|_default_manager|_base_manager)\.(\w+)\(\)(\[\])?$", hint)
+        if orm:
+            model = self.resolve_dotted(module, orm.group(1))
+            if model and model.kind == "def" and (orm.group(2) in INSTANCE_OPS or orm.group(3)):
+                return model
+            return None
+        called = hint.endswith("()")
+        target = hint[:-2] if called else hint
+        if target.endswith("()[]") or "(" in target:
+            return None
+        ref = self.resolve_dotted(module, target)  # module scope only: local hints never refer to locals
+        if ref is None or ref.kind != "def":
+            return ref if ref and ref.kind == "external" else None
+        d = self.defs[ref.key]
+        if d["kind"] == "class":
+            return ref
+        if called and d.get("returns"):
+            ret = self.resolve_dotted(d["module"], d["returns"][0])
+            if ret and ret.kind == "def" and self.defs[ret.key]["kind"] == "class":
+                return ret
+        return None
+
+    def member(self, ref: Ref, name: str) -> Ref | None:
+        if ref.kind == "module":
+            return self.resolve_name(ref.key, name)
+        if ref.kind == "external":
+            return Ref("external", f"{ref.key}.{name}")
+        if ref.kind == "var":
+            if ref.key in self._resolving:  # `app = app.with_options()` and similar self-references
+                return None
+            v = self.vars.get(ref.key)
+            value = v and v.get("value")
+            if value and value.get("t") == "call":
+                self._resolving.add(ref.key)
+                try:
+                    inst = self.instance_type(ref.module, f"{value['f']}()")
+                finally:
+                    self._resolving.discard(ref.key)
+                if inst is not None:
+                    return self.member(inst, name)
+            return None
+        d = self.defs.get(ref.key)
+        if d is None:
+            return None
+        if d["kind"] == "class":
+            found = self.find_in_class(ref.key, name)
+            if found:
+                return found
+            attr_type = self.class_attr_type(ref.key, name)
+            return attr_type
+        return None
+
+    # ---- classes ----
+
+    def enclosing_class(self, scope: dict[str, Any]) -> dict[str, Any] | None:
+        current: dict[str, Any] | None = scope
+        while current is not None:
+            if current["kind"] == "class":
+                return current
+            parent = current.get("parent")
+            current = self.defs.get(f"{current['module']}:{parent}") if parent else None
+        return None
+
+    def bases(self, class_key: str) -> list[Ref]:
+        d = self.defs.get(class_key)
+        if d is None:
+            return []
+        out = []
+        for base in d.get("bases", []):
+            ref = self.resolve_dotted(d["module"], base)
+            out.append(ref if ref else Ref("external", base))
+        return out
+
+    def mro(self, class_key: str) -> list[str]:
+        """The class and its project bases, depth first, without repeats."""
+        out: list[str] = []
+        stack = [class_key]
+        while stack:
+            key = stack.pop(0)
+            if key in out or key not in self.defs:
+                continue
+            out.append(key)
+            stack[0:0] = [b.key for b in self.bases(key) if b.kind == "def"]
+        return out
+
+    def external_bases(self, class_key: str) -> list[str]:
+        names: list[str] = []
+        for key in self.mro(class_key):
+            names += [b.key for b in self.bases(key) if b.kind == "external"]
+        return names
+
+    def find_in_class(self, class_key: str, name: str) -> Ref | None:
+        for key in self.mro(class_key):
+            child = self.children.get(key, {}).get(name)
+            if child:
+                return Ref("def", child)
+        return None
+
+    def class_attr_type(self, class_key: str, name: str) -> Ref | None:
+        for key in self.mro(class_key):
+            d = self.defs[key]
+            hint = d.get("self_attrs", {}).get(name)
+            if hint:
+                return self.instance_type(d["module"], hint)
+        return None
+
+    # ---- calls ----
+
+    def resolve_call(self, scope: dict[str, Any], callee: str) -> tuple[Ref | None, str]:
+        """Resolve a call made inside `scope`. Returns the target and a confidence word."""
+        module = scope["module"]
+        if callee.startswith("super()."):
+            cls = self.enclosing_class(scope)
+            name = bare(split_dotted(callee)[1]) if len(split_dotted(callee)) > 1 else ""
+            if cls and name:
+                for base in self.mro(cls["key"])[1:]:
+                    child = self.children.get(base, {}).get(name)
+                    if child:
+                        return Ref("def", child), "exact"
+            return None, "exact"
+        segments = split_dotted(callee)
+        if not segments:
+            return None, "exact"
+        if any("(" in s for s in segments[:-1]):  # a call on a call result: type unknown
+            return self._by_name(bare(segments[-1]))
+        ref = self.resolve_dotted(module, callee, scope)
+        if ref is not None:
+            return ref, "exact"
+        if len(segments) > 1:
+            return self._by_name(bare(segments[-1]))
+        return None, "exact"
+
+    def _by_name(self, name: str) -> tuple[Ref | None, str]:
+        if name in COMMON_METHODS or name.startswith("_"):
+            return None, "exact"
+        candidates = self.methods_by_name.get(name, [])
+        if len(candidates) == 1:
+            return Ref("def", candidates[0]), "inferred"
+        return None, "exact"
+
+    def const(self, module: str, value: dict[str, Any] | None) -> Any:
+        """Evaluate a literal-ish value: literals, names of module constants (followed through imports), lists,
+        concatenations of strings."""
+        if value is None:
+            return None
+        t = value.get("t")
+        if t == "lit":
+            return value["v"]
+        if t == "list":
+            return [self.const(module, v) for v in value["v"]]
+        if t == "concat":
+            parts = [self.const(module, v) for v in value["v"]]
+            return "".join(parts) if all(isinstance(p, str) for p in parts) else None
+        if t == "name":
+            ref = self.resolve_dotted(module, value["v"])
+            if ref and ref.kind == "var":
+                var = self.vars.get(ref.key)
+                if var and var.get("value") and var["value"].get("t") in ("lit", "concat", "list"):
+                    return self.const(ref.module, var["value"])
+        return None
+
+
+#: Django manager calls that return a model instance rather than a queryset.
+INSTANCE_OPS = frozenset(
+    {
+        "get",
+        "create",
+        "first",
+        "last",
+        "latest",
+        "earliest",
+        "get_or_create",
+        "update_or_create",
+        "aget",
+        "acreate",
+        "afirst",
+        "alast",
+    }
+)