fix: preserve runtime imports in TS type normalization

(cherry picked from commit ff7cb6d27a)

graphify/extract.py

@@ -1316,33 +1316,185 @@ def _extract_python_rationale(path: Path, result: dict) -> None:
 _TS_IMPORT_CALL_RE = re.compile(
     rb"\bimport\s*\(\s*['\"][^'\"\r\n]+['\"]\s*\)"
 )
-_TS_IMPORT_TYPE_CALL_RE = re.compile(
-    rb"<((?:[^;{}]*?\bimport\s*\([^()]+\)[^;{}]*?)+)>(?=\s*\()"
-)
 
 
-def _normalize_ts_import_types(source: bytes) -> bytes | None:
+def _ts_import_is_code(root: Any, start: int) -> bool:
-    """Rewrite TypeScript `import(...)` type arguments in call expressions
+    """Return whether an import-call match starts in executable source.
-    to standard type identifiers of identical byte length (#3154).
 
-    Preserves byte length, newlines, and source offsets so all downstream node
+    Regex matching is only used to locate a literal specifier; comments,
-    source_location metadata remains 100% accurate.
+    strings, regular expressions, and template text must never be fed into the
+    structural masking pass.  A template substitution is executable again, so
+    it is the one exception to the ``template_string`` guard.
     """
-    if rb"import(" not in source and rb"import (" not in source:
+    node = root.descendant_for_byte_range(start, start + 1)
+    if node is None:
+        # A malformed tree can leave a byte range uncovered. Treat it as code
+        # and let the structural pass decide whether it belongs to a type list;
+        # never turn an uncertain lexical result into a crash.
+        return True
+    in_template_substitution = False
+    while node is not None:
+        if node.type == "comment" or node.type in ("string", "regex", "regex_pattern"):
+            return False
+        if node.type == "template_substitution":
+            in_template_substitution = True
+        elif node.type == "template_string" and not in_template_substitution:
+            return False
+        node = node.parent
+    return True
+
+
+def _ts_type_argument_ranges(root: Any, *, call_only: bool) -> list[tuple[int, int]]:
+    """Collect byte ranges tree-sitter already parsed as type arguments."""
+    ranges: list[tuple[int, int]] = []
+    stack = [root]
+    while stack:
+        node = stack.pop()
+        if node.type == "type_arguments":
+            parent = node.parent
+            if not call_only or (
+                parent is not None and parent.type in ("call_expression", "new_expression")
+            ):
+                ranges.append((node.start_byte, node.end_byte))
+        stack.extend(node.children)
+    return ranges
+
+
+def _ts_error_nodes(root: Any) -> list[Any]:
+    """Return parser error nodes without depending on a grammar's error name."""
+    errors: list[Any] = []
+    stack = [root]
+    while stack:
+        node = stack.pop()
+        if node.type == "ERROR" or node.is_error:
+            errors.append(node)
+        stack.extend(node.children)
+    return errors
+
+
+def _ts_mask_candidate_is_malformed(
+    original_root: Any,
+    masked_range: tuple[int, int],
+    errors: list[Any],
+) -> bool:
+    """Tell whether a masked generic is backed by an actual parse failure.
+
+    A valid runtime comparison can have the same token shape as a generic call
+    after the import is replaced (``a < import("x") > (a)``).  Tree-sitter
+    exposes the opening ``<`` as a binary operator in the original tree for
+    both forms, so the structural second pass alone cannot distinguish them.
+    Only repair that ambiguity when the original parser has an error adjacent
+    to the candidate's closing angle; that is the signature of the known
+    ``import(...)``-type grammar failure.  Valid comparisons, including ones
+    nested in another call's arguments, remain byte-for-byte untouched.
+    """
+    start, end = masked_range
+    opener = original_root.descendant_for_byte_range(start, start + 1)
+    if opener is None or opener.type != "<":
+        # A grammar that already gives us a structural type_arguments node is
+        # safe to normalize; no binary/comparison ambiguity is present.
+        return True
+    if opener.parent is None or opener.parent.type != "binary_expression":
+        return True
+
+    # A malformed generic's ERROR starts at the closing `>` (or at the call
+    # punctuation immediately following it). Keep this window deliberately
+    # narrow so an unrelated syntax error elsewhere cannot authorize masking a
+    # valid runtime comparison.
+    for error in errors:
+        if error.start_byte <= end + 2 and error.end_byte >= end - 1:
+            return True
+    return False
+
+
+def _normalize_ts_import_types(source: bytes, *, tsx: bool = False) -> bytes | None:
+    """Rewrite only syntactic TypeScript ``import(...)`` type arguments.
+
+    The first implementation of #3154 used a ``<...>`` regular expression.
+    In semicolon-less code that expression could span two comparison
+    operators, so a *runtime* dynamic import was blanked before parsing.  A
+    temporary, byte-preserving mask lets tree-sitter identify the actual
+    ``type_arguments`` node without asking it to parse the known-invalid
+    ``import(...)`` call-site form.  We then rewrite only placeholders inside
+    call/new-expression type arguments.  Keeping every replacement the same
+    byte length preserves source offsets and line locations.
+    """
+    raw_matches = list(_TS_IMPORT_CALL_RE.finditer(source))
+    if not raw_matches:
         return None
 
-    def repl_type_args(m: "re.Match[bytes]") -> bytes:
+    # tree-sitter is already a required TypeScript dependency for extraction.
-        type_arg_content = m.group(1)
+    # If it cannot be loaded here, leave the source untouched; the normal AST
+    # path will report its own dependency/parser error rather than applying an
+    # unsafe textual guess.
+    try:
+        import tree_sitter_typescript as ts_typescript
+        from tree_sitter import Language, Parser
 
-        def repl_import(im: "re.Match[bytes]") -> bytes:
+        language_factory = (
-            matched = im.group(0)
+            ts_typescript.language_tsx if tsx else ts_typescript.language_typescript
-            return b"T" + re.sub(rb"[^\r\n]", b" ", matched[1:])
+        )
+        parser = Parser(Language(language_factory()))
+        original_root = parser.parse(source).root_node
+    except Exception:
+        return None
 
-        new_content = _TS_IMPORT_CALL_RE.sub(repl_import, type_arg_content)
+    matches = [
-        return b"<" + new_content + b">"
+        match for match in raw_matches
+        if _ts_import_is_code(original_root, match.start())
+    ]
+    if not matches:
+        return None
 
-    norm = _TS_IMPORT_TYPE_CALL_RE.sub(repl_type_args, source)
+    # If the original tree already has a type_arguments node around a match,
-    return norm if norm != source else None
+    # its syntax is parseable as written. This includes the grammar's deliberate
+    # comparison-vs-generic ambiguity (`a < b, import("...") > (d)`); retaining
+    # that source is essential because it is a runtime expression, not a type.
+    original_type_ranges = _ts_type_argument_ranges(original_root, call_only=False)
+    matches = [
+        match for match in matches
+        if not any(start <= match.start() < end for start, end in original_type_ranges)
+    ]
+    if not matches:
+        return None
+
+    def placeholder(match: "re.Match[bytes]") -> bytes:
+        # Keep CR/LF bytes intact. The first byte becomes an ordinary type
+        # identifier and the remaining bytes are padding, so every downstream
+        # byte offset remains identical to the user's source.
+        matched = match.group(0)
+        return b"T" + re.sub(rb"[^\r\n]", b" ", matched[1:])
+
+    masked = bytearray(source)
+    for match in matches:
+        masked[match.start():match.end()] = placeholder(match)
+
+    root = parser.parse(bytes(masked)).root_node
+
+    # Only call/new-expression type arguments need this workaround. Ordinary
+    # type annotations already parse ``import(...)`` correctly and should keep
+    # their native AST shape. A range from an outer call's type_arguments also
+    # covers nested generic arguments.
+    type_argument_ranges = _ts_type_argument_ranges(root, call_only=True)
+    errors = _ts_error_nodes(original_root)
+
+    if not type_argument_ranges:
+        return None
+
+    norm = bytearray(source)
+    changed = False
+    for match in matches:
+        containing_ranges = [
+            candidate for candidate in type_argument_ranges
+            if candidate[0] <= match.start() < candidate[1]
+        ]
+        if containing_ranges and any(
+            _ts_mask_candidate_is_malformed(original_root, candidate, errors)
+            for candidate in containing_ranges
+        ):
+            norm[match.start():match.end()] = placeholder(match)
+            changed = True
+    return bytes(norm) if changed else None
 
 
 # ── Public API ────────────────────────────────────────────────────────────────
@@ -1369,7 +1521,7 @@ def extract_js(path: Path) -> dict:
     if is_ts:
         try:
             source = path.read_bytes()
-            source_override = _normalize_ts_import_types(source)
+            source_override = _normalize_ts_import_types(source, tsx=suffix == ".tsx")
         except OSError:
             pass
     result = _extract_generic(path, config, source_override=source_override)
@@ -1407,7 +1559,7 @@ def _rescue_js_dynamic_imports(path: Path, result: dict) -> None:
     try:
         import re as _re
         src = path.read_text(encoding="utf-8", errors="replace")
-        if "import(" not in src:  # cheap bail — most files have none
+        if not _re.search(r"(?<!\w)import\s*\(", src):  # cheap bail — most files have none
             return
         existing_ids = {n["id"] for n in result.get("nodes", [])}
         file_node_id = _make_id(str(path))
@@ -1449,7 +1601,7 @@ def _rescue_js_dynamic_imports(path: Path, result: dict) -> None:
         # handling: a literal `import(`./x`)` resolves, `${`-substituted ones
         # are excluded (no `$` in the class) as statically unresolvable.
         for m in _re.finditer(
-            r"""(?<!\w)import\(\s*(?:'([^'\n]+)'|"([^"\n]+)"|`([^`$\n]+)`)\s*\)""",
+            r"""(?<!\w)import\s*\(\s*(?:'([^'\n]+)'|"([^"\n]+)"|`([^`$\n]+)`)\s*\)""",
             src,
         ):
             raw = m.group(1) or m.group(2) or m.group(3)
@@ -1839,7 +1991,9 @@ def extract_vue(path: Path) -> dict:
         config = _TS_CONFIG
     masked_bytes = masked.encode("utf-8")
     if config in (_TS_CONFIG, _TSX_CONFIG):
-        masked_bytes = _normalize_ts_import_types(masked_bytes) or masked_bytes
+        masked_bytes = _normalize_ts_import_types(
+            masked_bytes, tsx=config is _TSX_CONFIG
+        ) or masked_bytes
 
     result = _extract_generic(path, config, source_override=masked_bytes)
 

tests/test_ts_import_type_arguments.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 import os
 from pathlib import Path
 
-from graphify.extract import extract
+from graphify.extract import _normalize_ts_import_types, extract
 
 
 def _extract(tmp_path: Path, files: dict[str, str]):
@@ -34,6 +34,15 @@ def _labels(r: dict) -> set[str]:
     return {n["label"] for n in r["nodes"]}
 
 
+def _labelled_edges(r: dict) -> set[tuple[str, str, str]]:
+    labels = {n["id"]: n["label"] for n in r["nodes"]}
+    return {
+        (labels.get(e["source"], e["source"]), e["relation"],
+         labels.get(e["target"], e["target"]))
+        for e in r["edges"]
+    }
+
+
 def _assert_silent(err: str):
     assert "syntax errors" not in err
     assert "partially extracted" not in err
@@ -139,3 +148,140 @@ def test_ts_multiline_import_type_arguments(tmp_path: Path, capsys):
     assert "after()" in labels
     assert "mod" in labels
     _assert_silent(capsys.readouterr().err)
+
+
+def test_ts_runtime_dynamic_import_between_comparisons_is_not_normalized(tmp_path: Path):
+    """#3210: ``<`` and ``>`` in separate semicolon-less statements must not
+    make a runtime import look like a call-expression type argument."""
+    r = _extract(tmp_path, {
+        "main.ts": (
+            "async function load(a: number, b: number) {\n"
+            "  const flag = a < b\n"
+            "  const m = await import('./mod')\n"
+            "  const ok = b > (a)\n"
+            "  return [flag, m, ok]\n"
+            "}\n"
+        ),
+        "mod.ts": "export const value = 1\n",
+    })
+    edges = _labelled_edges(r)
+    assert ("load()", "imports_from", "mod.ts") in edges
+    assert ("main.ts", "dynamic_import", "mod.ts") in edges
+
+
+def test_ts_spaced_runtime_dynamic_import_keeps_both_edge_granularities(tmp_path: Path):
+    """The normalizer and rescue both accept whitespace before ``(``."""
+    r = _extract(tmp_path, {
+        "main.ts": (
+            "async function load(a: number, b: number) {\n"
+            "  const flag = a < b\n"
+            "  const m = await import ('./mod')\n"
+            "  const ok = b > (a)\n"
+            "  return [flag, m, ok]\n"
+            "}\n"
+        ),
+        "mod.ts": "export const value = 1\n",
+    })
+    edges = _labelled_edges(r)
+    assert ("load()", "imports_from", "mod.ts") in edges
+    assert ("main.ts", "dynamic_import", "mod.ts") in edges
+
+
+def test_ts_multiple_runtime_imports_between_comparisons_survive(tmp_path: Path):
+    r = _extract(tmp_path, {
+        "main.ts": (
+            "async function load(a: number, b: number) {\n"
+            "  const flag = a < b\n"
+            "  const m = await import('./mod')\n"
+            "  const o = await import ('./other')\n"
+            "  const ok = b > (a)\n"
+            "  return [flag, m, o, ok]\n"
+            "}\n"
+        ),
+        "mod.ts": "export const value = 1\n",
+        "other.ts": "export const other = 2\n",
+    })
+    edges = _labelled_edges(r)
+    for target in ("mod.ts", "other.ts"):
+        assert ("load()", "imports_from", target) in edges
+        assert ("main.ts", "dynamic_import", target) in edges
+
+
+def test_ts_nested_comparison_runtime_imports_are_not_masked(tmp_path: Path):
+    """A comparison expression inside a call argument is not a generic call.
+
+    tree-sitter can represent ``a < b, import('./mod') > (d)`` as a nested
+    call with ``type_arguments`` after a placeholder is inserted. The original
+    tree already parses this runtime expression correctly, so it must remain
+    untouched in both named-function and arrow-function bodies.
+    """
+    sources = (
+        "async function load(a: number, b: number, d: number) {\n"
+        "  const value = foo(a < b, import('./mod') > (d))\n"
+        "  return value\n"
+        "}\n",
+        "const load = async (a: number, b: number, d: number) => {\n"
+        "  const value = foo(a < b, import  ('./mod') > (d))\n"
+        "  return value\n"
+        "}\n",
+    )
+    for index, source in enumerate(sources):
+        case = tmp_path / f"case-{index}"
+        r = _extract(case, {
+            "main.ts": source,
+            "mod.ts": "export const value = 1\n",
+        })
+        edges = _labelled_edges(r)
+        assert ("load()", "imports_from", "mod.ts") in edges
+        assert ("main.ts", "dynamic_import", "mod.ts") in edges
+
+
+def test_ts_comparison_with_import_as_middle_operand_is_not_masked(tmp_path: Path):
+    """A valid ``a < import('mod') > (a)`` chain is runtime code, not a type."""
+    sources = (
+        "function load(a: number) {\n"
+        "  const value = a < import('./mod') > (a)\n"
+        "  return value\n"
+        "}\n",
+        "function load(a: number) {\n"
+        "  const value = a < import  ('./mod') > (a)\n"
+        "  return value\n"
+        "}\n",
+    )
+    for index, source in enumerate(sources):
+        case = tmp_path / f"case-{index}"
+        r = _extract(case, {
+            "main.ts": source,
+            "mod.ts": "export const value = 1\n",
+        })
+        edges = _labelled_edges(r)
+        assert ("load()", "imports_from", "mod.ts") in edges
+        assert ("main.ts", "dynamic_import", "mod.ts") in edges
+
+
+def test_ts_normalizer_masks_only_structural_call_type_arguments():
+    source = (
+        b"const flag = a < b\n"
+        b"const m = await import ('./runtime')\n"
+        b"const ok = b > (a)\n"
+        b"load<typeof import('./types')>();\n"
+    )
+    normalized = _normalize_ts_import_types(source)
+    assert normalized is not None
+    assert b"import ('./runtime')" in normalized
+    assert b"import('./types')" not in normalized
+    assert len(normalized) == len(source)
+    assert normalized.count(b"\n") == source.count(b"\n")
+
+
+def test_ts_normalizer_does_not_mask_import_text_in_literals_or_comments():
+    source = (
+        b'const text = "load<typeof import(\\\"./text\\\")>()";\n'
+        b'// load<typeof import("./comment")>();\n'
+        b'load<typeof import("./types")>();\n'
+    )
+    normalized = _normalize_ts_import_types(source)
+    assert normalized is not None
+    assert b'import(\\\"./text\\\")' in normalized
+    assert b'import("./comment")' in normalized
+    assert b'import("./types")' not in normalized