From 291248ef49577205026effd144cc2b1eafc42e2b Mon Sep 17 00:00:00 2001 From: madhusudanpatnaik Date: Fri, 23 Oct 2026 18:15:06 +0100 Subject: [PATCH] fix(python): extract explicit type aliases as type nodes (#4578) Applied from PR #4578 (authorship preserved); the PR's stale cache-schema bump is dropped in favour of one combined bump for the alias cluster. Co-Authored-By: Claude Opus 4.8 --- graphify/extractors/engine.py | 63 ++++++++++++++- tests/test_python_type_alias.py | 135 ++++++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 tests/test_python_type_alias.py diff --git a/graphify/extractors/engine.py b/graphify/extractors/engine.py index cc9fbae87..6e87e70a0 100644 --- a/graphify/extractors/engine.py +++ b/graphify/extractors/engine.py @@ -1235,6 +1235,49 @@ def _swift_declaration_keyword(node) -> str | None: return c.type return None +_PY_TYPE_ALIAS_ANNOTATIONS = frozenset({ + "TypeAlias", "typing.TypeAlias", "typing_extensions.TypeAlias", +}) +_PY_NEW_TYPE_CALLEES = frozenset({"NewType", "typing.NewType", "typing_extensions.NewType"}) + + +def _python_alias_name(stmt, source: bytes) -> str | None: + """The type name a module-level statement declares as an explicit alias. + + `type Vector = ...` / `type Pair[T] = ...` (PEP 695), `Matrix: TypeAlias = ...` + and `UserId = NewType("UserId", int)` each declare one. A bare + `JsonValue = Union[...]` is an alias by convention only and is not one. + A same-name re-export (`Field: TypeAlias = core.Field`, `type Field = + core.Field[int]`, or `Field: TypeAlias = Field` after `from core import + Field`) denotes the original type and declares no new name, so it + returns None: references keep resolving to the one real `Field`. + """ + if stmt.type == "type_alias_statement": + left, right = stmt.child_by_field_name("left"), stmt.child_by_field_name("right") + name = _read_text(left, source).split("[", 1)[0].strip() if left is not None else "" + elif stmt.type == "expression_statement" and len(stmt.named_children) == 1: + assign = stmt.named_children[0] + left = assign.child_by_field_name("left") if assign.type == "assignment" else None + if left is None or left.type != "identifier": + return None + name, right = _read_text(left, source), assign.child_by_field_name("right") + annotation = assign.child_by_field_name("type") + if annotation is None: + callee = right.child_by_field_name("function") if right is not None and right.type == "call" else None + # NewType always mints a new type, even under its own name. + return name if callee is not None and _read_text(callee, source) in _PY_NEW_TYPE_CALLEES else None + if _read_text(annotation, source).strip() not in _PY_TYPE_ALIAS_ANNOTATIONS: + return None + else: + return None + if not name.isidentifier(): + return None + target = _read_text(right, source).split("[", 1)[0].strip() if right is not None else "" + if target.rsplit(".", 1)[-1] == name: + return None + return name + + def _python_pre_scan_underscore_collisions(root_node, source: bytes, stem: str) -> dict[str, set[str]]: """Pre-scan a Python module for name-only differences that collapse to one node id. @@ -1275,7 +1318,10 @@ def _python_pre_scan_underscore_collisions(root_node, source: bytes, stem: str) classes: list = [] # (class_name, body) for the method pass below for child in map(_undecorated, root_node.children): - if child.type == "function_definition": + alias_name = _python_alias_name(child, source) + if alias_name: + _record(_make_id(stem, alias_name), alias_name) + elif child.type == "function_definition": name_node = child.child_by_field_name("name") if name_node is not None: name = _read_text(name_node, source) @@ -8299,6 +8345,21 @@ def _extract_generic( ruby_namespace, ruby_lexical_scopes): return + # An explicit module-level type alias declares a type name: give it a + # node so annotations using it reference it instead of minting a + # sourceless stub. The statement is then walked as usual. + if (t in ("type_alias_statement", "expression_statement") + and config.ts_module == "tree_sitter_python" and parent_class_nid is None + and node.parent is not None and node.parent.type == "module"): + alias_name = _python_alias_name(node, source) + if alias_name: + alias_line = node.start_point[0] + 1 + alias_nid = _underscore_salted_nid( + _make_id(stem, alias_name), alias_name, python_underscore_groups + ) + add_node(alias_nid, alias_name, alias_line) + add_edge(file_nid, alias_nid, "contains", alias_line) + # Python's `@property` / `@staticmethod` / `@classmethod` wrap the # inner function_definition in a `decorated_definition` node. The # default recurse below clears parent_class_nid, which would cause the diff --git a/tests/test_python_type_alias.py b/tests/test_python_type_alias.py new file mode 100644 index 000000000..b9e9f1ac4 --- /dev/null +++ b/tests/test_python_type_alias.py @@ -0,0 +1,135 @@ +"""Explicit Python type aliases are declarations. + +A module-level `type Vector = list[float]` (PEP 695), `Matrix: TypeAlias = ...` +and `UserId = NewType("UserId", int)` each declare a type name, but the Python +extractor emitted no node for them. Every annotation using the name, even in +the same module, resolved to a SOURCELESS stub that the corpus rewire can grow +into a cross-file hub. They are now type nodes contained by their module, in +the same collision census as module-level classes and functions. + +Boundaries: only these explicit forms count. A bare `JsonValue = Union[...]` +is an alias by convention only, and treating every capitalised assignment as a +type would be a guess. A same-name re-export (`Field: TypeAlias = core.Field`, +`type Field = core.Field[int]`, `Field: TypeAlias = Field` after +`from core import Field`) denotes the original type and declares no new +name, so it emits no node: references keep resolving to the one real `Field`. +""" +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest + +from graphify.extract import extract + +_SOURCE = textwrap.dedent("""\ + import typing + from typing import NewType, TypeAlias + + type Vector = list[float] + type Pair[T] = tuple[T, T] + UserId = NewType("UserId", int) + Matrix: TypeAlias = list[Vector] + Grid: typing.TypeAlias = list[Matrix] + + + def scale(v: Vector, u: UserId, m: Matrix, p: Pair[int], g: Grid) -> Vector: + return v + """) + + +def _extract(tmp_path: Path, files: dict[str, str]) -> dict: + for name, body in files.items(): + (tmp_path / name).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / name).write_text(body, encoding="utf-8") + return extract([tmp_path / n for n in files], root=tmp_path, + cache_root=tmp_path / "graphify-out", parallel=False) + + +def _labels(result: dict) -> dict[str, str]: + return {n["id"]: n["label"] for n in result["nodes"]} + + +def _pairs(result: dict, relation: str) -> set[tuple[str, str]]: + label = _labels(result) + return {(label.get(e["source"], str(e["source"])), label.get(e["target"], str(e["target"]))) + for e in result["edges"] if e["relation"] == relation} + + +def test_each_explicit_alias_is_a_sourced_node_on_its_line(tmp_path: Path): + result = _extract(tmp_path, {"shapes.py": _SOURCE}) + for label, line in (("Vector", "L4"), ("Pair", "L5"), ("UserId", "L6"), + ("Matrix", "L7"), ("Grid", "L8")): + matches = [n for n in result["nodes"] if n["label"] == label] + assert len(matches) == 1, label + assert matches[0]["source_file"], label + assert matches[0]["source_location"] == line, label + assert [n["label"] for n in result["nodes"] if not n.get("source_file")] == [] + + +def test_each_alias_is_contained_by_its_module(tmp_path: Path): + contained = {tgt for _src, tgt in _pairs(_extract(tmp_path, {"shapes.py": _SOURCE}), "contains")} + assert {"Vector", "Pair", "UserId", "Matrix", "Grid"} <= contained + + +def test_annotations_reference_the_local_aliases(tmp_path: Path): + refs = _pairs(_extract(tmp_path, {"shapes.py": _SOURCE}), "references") + assert {("scale()", "Vector"), ("scale()", "UserId"), ("scale()", "Matrix"), + ("scale()", "Pair"), ("scale()", "Grid")} <= refs + + +def test_an_alias_and_a_function_that_fold_to_one_id_stay_separate(tmp_path: Path): + """`Vector` and `vector()` slug to one id (#2810's census).""" + result = _extract(tmp_path, {"shapes.py": "type Vector = list[float]\n\n\ndef vector():\n pass\n"}) + labels = sorted(n["label"] for n in result["nodes"] if n["label"] in ("Vector", "vector()")) + assert labels == ["Vector", "vector()"] + assert len({n["id"] for n in result["nodes"] if n["label"] in ("Vector", "vector()")}) == 2 + + +def test_a_reference_from_another_module_lands_on_the_alias(tmp_path: Path): + result = _extract(tmp_path, { + "types_.py": "from typing import NewType\n\nUserId = NewType(\"UserId\", int)\n", + "users.py": "from types_ import UserId\n\n\ndef load(uid: UserId) -> None:\n pass\n", + }) + [uid] = [n for n in result["nodes"] if n["label"] == "UserId"] + assert uid["source_file"].endswith("types_.py") + label = _labels(result) + targets = {e["target"] for e in result["edges"] + if e["relation"] == "references" and label.get(e["source"]) == "load()"} + assert uid["id"] in targets + + +@pytest.mark.parametrize("alias", [ + "from typing import TypeAlias\nimport core\n\nField: TypeAlias = core.Field\n", + "import core\n\ntype Field = core.Field[int]\n", +]) +def test_a_same_name_re_export_emits_no_node(tmp_path: Path, alias: str): + result = _extract(tmp_path, {"api.py": alias}) + assert [n for n in result["nodes"] if n["label"] == "Field"] == [] + + + +def test_a_bare_same_name_re_export_keeps_references_on_the_original(tmp_path: Path): + result = _extract(tmp_path, { + "core.py": "class Field:\n pass\n", + "api.py": ("from typing import TypeAlias\nfrom core import Field\n\n" + "Field: TypeAlias = Field\n\ndef f(x: Field) -> None: ...\n"), + }) + [field] = [n for n in result["nodes"] if n["label"] == "Field"] + assert field["source_file"].endswith("core.py") + label = _labels(result) + targets = {e["target"] for e in result["edges"] + if e["relation"] == "references" and label.get(e["source"]) == "f()"} + assert targets == {field["id"]} + +@pytest.mark.parametrize("statement", [ + "JsonValue = dict[str, object]\n", # an alias by convention only + "limit: int = 3\n", # an annotated variable + "def f():\n type Local = int\n", # not module level + "class C:\n Inner: TypeAlias = int\n", # not module level +]) +def test_only_explicit_module_level_aliases_are_declarations(tmp_path: Path, statement: str): + result = _extract(tmp_path, {"m.py": "from typing import TypeAlias\n\n" + statement}) + names = {"JsonValue", "limit", "Local", "Inner"} + assert [n["label"] for n in result["nodes"] if n["label"] in names] == []