fix(extract): stop attribute-form Python decorators binding to an unrelated def (#2315)

graphify/extractors/engine.py

@@ -2417,6 +2417,15 @@ def ensure_named_node(name: str, line: int) -> str:
     file_nid = _make_id(str(path))
     add_node(file_nid, path.name, 1)
 
+    # Attribute-form decorator receivers (#2315): computed once so the
+    # decorated_definition branch below can tell a locally-assigned receiver
+    # apart from an imported module/symbol. See _python_local_instance_classes.
+    local_assigned_names: set[str] = set()
+    local_receiver_classes: dict[str, str] = {}
+    if config.ts_module == "tree_sitter_python":
+        local_assigned_names, local_receiver_classes = (
+            _python_local_instance_classes(root, source, config.class_types))
+
     def walk(node, parent_class_nid: str | None = None) -> None:
         t = node.type
 
@@ -3781,7 +3790,33 @@ def _emit_java_parent_type(type_node, rel: str, at_line: int) -> None:
                         if not deco_name or deco_name in _PYTHON_DECORATOR_NOISE:
                             continue
                         deco_line = child.start_point[0] + 1
-                        target = ensure_named_node(deco_name, deco_line)
+                        # Attribute-form receiver (#2315): a bare-name lookup on
+                        # just the symbol (`get` in `app.get(...)`) is free to
+                        # collide with any unrelated same-named def, same-file or
+                        # via the corpus rewire. A receiver bound to a locally-
+                        # instantiated local class targets that class's method
+                        # directly; a receiver that's locally assigned but not
+                        # traced to a class (or a compound `@a.b.c` chain)
+                        # gets a receiver-scoped stub instead of the bare name so
+                        # it can never bind to an unrelated definition. A receiver
+                        # that's neither (i.e. likely imported) keeps the existing
+                        # bare-name behavior locked by
+                        # test_attribute_decorator_targets_the_symbol_not_the_module.
+                        receiver = _python_decorator_receiver(child, source)
+                        if receiver is None:
+                            target = ensure_named_node(deco_name, deco_line)
+                        elif receiver in local_receiver_classes:
+                            class_nid = _make_id(
+                                stem, "", local_receiver_classes[receiver])
+                            target = _make_id(class_nid, deco_name)
+                        elif "." in receiver or receiver in local_assigned_names:
+                            # A compound chain (`a.b`) can't be a plain import
+                            # alias either, so it gets the same scoped-stub
+                            # treatment as an unresolved local variable.
+                            target = ensure_named_node(
+                                f"{receiver}.{deco_name}", deco_line)
+                        else:
+                            target = ensure_named_node(deco_name, deco_line)
                         if target != owner_nid:
                             add_edge(owner_nid, target, "references", deco_line,
                                      context="decorator")
@@ -4744,6 +4779,67 @@ def _python_decorator_name(deco_node, source: bytes) -> str | None:
         return None
     return None
 
+def _python_decorator_receiver(deco_node, source: bytes) -> str | None:
+    """Return the receiver text of an attribute-form Python decorator.
+
+    `"app"` for `@app.get(...)` / `@app.route(...)`; the full dotted text
+    (`"a.b"`) when the decorator IS attribute-form but its object isn't a
+    plain identifier (`@a.b.c`, an unresolvable chain); `None` when the
+    decorator isn't attribute-form at all (`@traced`, `@retry(...)`) — those
+    have no receiver-collision risk (#2315).
+    """
+    for child in deco_node.children:
+        if not child.is_named:
+            continue
+        target = child
+        if target.type == "call":
+            target = target.child_by_field_name("function") or target
+        if target.type == "attribute":
+            obj = target.child_by_field_name("object")
+            return _read_text(obj, source) if obj else None
+        return None
+    return None
+
+def _python_local_instance_classes(
+    root, source: bytes, class_types
+) -> tuple[set[str], dict[str, str]]:
+    """Module-level `name = expr` targets, and the subset assigned a call to a
+    class defined in this same file (`name = ClassName()`) mapped to the class
+    name (#2315).
+
+    Scoped to simple top-level `name = call(...)` assignments — no control
+    flow, no attribute targets — matching the shape of `app = _App()` without
+    building general type inference. Used so an attribute-form decorator's
+    receiver (`app` in `@app.get(...)`) can be told apart from an imported
+    module/symbol, which never appears on the left of a plain assignment.
+    """
+    local_class_names = {
+        _read_text(n.child_by_field_name("name"), source)
+        for n in root.children
+        if n.type in class_types and n.child_by_field_name("name")
+    }
+    assigned: set[str] = set()
+    receiver_classes: dict[str, str] = {}
+    for node in root.children:
+        stmt = node
+        if stmt.type == "expression_statement" and stmt.named_child_count == 1:
+            stmt = stmt.named_children[0]
+        if stmt.type != "assignment":
+            continue
+        left = stmt.child_by_field_name("left")
+        right = stmt.child_by_field_name("right")
+        if left is None or left.type != "identifier":
+            continue
+        left_name = _read_text(left, source)
+        assigned.add(left_name)
+        if right is not None and right.type == "call":
+            fn = right.child_by_field_name("function")
+            if fn is not None and fn.type == "identifier":
+                cls_name = _read_text(fn, source)
+                if cls_name in local_class_names:
+                    receiver_classes[left_name] = cls_name
+    return assigned, receiver_classes
+
 def _ts_decorator_name(deco_node, source: bytes) -> str | None:
     """Return the head symbol of a TS `decorator` node.
 

tests/test_python_decorators.py

@@ -225,3 +225,97 @@ def test_undecorated_function_emits_no_decorator_edge(tmp_path):
                "    pass\n")
     r = extract([f], cache_root=tmp_path)
     assert _deco_edges(r, _func_nid("pkg/plain.py", "plain")) == set()
+
+
+def test_attribute_decorator_local_receiver_targets_class_method(tmp_path):
+    # `app = _App()` then `@app.get(...)`: the receiver is a locally-
+    # instantiated local class, so the edge must land on `_App.get`, not a
+    # bare "get" stub free to collide with an unrelated same-named def
+    # elsewhere in the corpus (#2315).
+    routes = _write(tmp_path / "pkg" / "routes.py",
+                     "class _App:\n"
+                     "    def get(self, path):\n"
+                     "        pass\n"
+                     "\n"
+                     "app = _App()\n"
+                     "\n"
+                     "@app.get(\"/health\")\n"
+                     "def health():\n"
+                     "    pass\n"
+                     "\n"
+                     "@app.get(\"/status\")\n"
+                     "def status():\n"
+                     "    pass\n")
+    unrelated = _write(tmp_path / "pkg" / "settings_store.py",
+                        "def get(key):\n"
+                        "    pass\n")
+    r = extract([routes, unrelated], cache_root=tmp_path)
+    unrelated_nid = _func_nid("pkg/settings_store.py", "get")
+    method_nid = _method_nid("pkg/routes.py", "_App", "get")
+    # The edge must land on a node that actually exists, not a dangling id.
+    assert method_nid in {n["id"] for n in r["nodes"]}
+    for func in ("health", "status"):
+        targets = _deco_edges(r, _func_nid("pkg/routes.py", func))
+        assert unrelated_nid not in targets
+        assert method_nid in targets
+
+
+def test_attribute_decorator_local_receiver_same_file_unrelated_def(tmp_path):
+    # Same collision shape, but the unrelated `def get` lives in the SAME
+    # file as the decorated function, at the same (module) scope, which the
+    # scoped-id fast path in ensure_named_node would otherwise match directly.
+    f = _write(tmp_path / "pkg" / "mixed.py",
+               "class _App:\n"
+               "    def get(self, path):\n"
+               "        pass\n"
+               "\n"
+               "app = _App()\n"
+               "\n"
+               "def get(key):\n"
+               "    pass\n"
+               "\n"
+               "@app.get(\"/health\")\n"
+               "def health():\n"
+               "    pass\n")
+    r = extract([f], cache_root=tmp_path)
+    unrelated_nid = _func_nid("pkg/mixed.py", "get")
+    method_nid = _method_nid("pkg/mixed.py", "_App", "get")
+    # The edge must land on a node that actually exists, not a dangling id.
+    assert method_nid in {n["id"] for n in r["nodes"]}
+    targets = _deco_edges(r, _func_nid("pkg/mixed.py", "health"))
+    assert unrelated_nid not in targets
+    assert method_nid in targets
+
+
+def test_bare_decorator_still_emits_edge(tmp_path):
+    # Regression guard for #2154: a plain identifier decorator is unaffected
+    # by the attribute-receiver handling added for #2315.
+    f = _write(tmp_path / "pkg" / "bare.py",
+               "def my_decorator(fn):\n"
+               "    return fn\n"
+               "\n"
+               "@my_decorator\n"
+               "def business_logic():\n"
+               "    pass\n")
+    r = extract([f], cache_root=tmp_path)
+    assert _func_nid("pkg/bare.py", "my_decorator") in _deco_edges(
+        r, _func_nid("pkg/bare.py", "business_logic"))
+
+
+def test_deep_attribute_decorator_does_not_bind_unrelated_def(tmp_path):
+    # `@a.b.c(...)`: the receiver is a compound attribute chain, not a plain
+    # identifier, so it can't be resolved to a local class either. It must not
+    # bind to an unrelated module-level `def c` elsewhere in the corpus.
+    deep = _write(tmp_path / "pkg" / "deep.py",
+                  "import a\n"
+                  "\n"
+                  "@a.b.c(\"/x\")\n"
+                  "def handler():\n"
+                  "    pass\n")
+    unrelated = _write(tmp_path / "pkg" / "other.py",
+                        "def c():\n"
+                        "    pass\n")
+    r = extract([deep, unrelated], cache_root=tmp_path)
+    unrelated_nid = _func_nid("pkg/other.py", "c")
+    targets = _deco_edges(r, _func_nid("pkg/deep.py", "handler"))
+    assert unrelated_nid not in targets