fixup! test(core.utils): use tokenize to parse SUPPORTED_TARGETS (ignore string/comment braces)

Author: cpcloud

cuda_core/tests/test_program_cache.py

@@ -270,44 +270,55 @@ def test_make_program_cache_key_rejects(kwargs, exc_type, match):
 
 def test_make_program_cache_key_supported_targets_matches_program_compile():
     """``_SUPPORTED_TARGETS_BY_CODE_TYPE`` duplicates the backend target
-    matrix in ``_program.pyx``. Guard against drift: parse the pyx source,
-    extract its SUPPORTED_TARGETS dict, and assert the two views agree for
-    every code_type."""
+    matrix in ``_program.pyx``. Guard against drift: parse the pyx source
+    with :mod:`tokenize` (which skips string literals and comments) to
+    extract ``SUPPORTED_TARGETS`` and assert the two views agree."""
     import ast
+    import io
+    import tokenize
     from pathlib import Path
 
     from cuda.core.utils._program_cache import _SUPPORTED_TARGETS_BY_CODE_TYPE
 
-    # Map Program._backend strings to the public code_type values.
     backend_to_code_type = {"NVRTC": "c++", "NVVM": "nvvm"}
     linker_backends = ("nvJitLink", "driver")
 
     pyx = Path(__file__).parent.parent / "cuda" / "core" / "_program.pyx"
     text = pyx.read_text()
-    # Find the assignment: ``cdef dict SUPPORTED_TARGETS = { ... }``.
-    marker = "cdef dict SUPPORTED_TARGETS"
-    start = text.index(marker)
-    brace = text.index("{", start)
-    # Find the matching closing brace.
+    marker_idx = text.index("cdef dict SUPPORTED_TARGETS")
+    tokens = tokenize.generate_tokens(io.StringIO(text[marker_idx:]).readline)
+
     depth = 0
-    for idx in range(brace, len(text)):
-        if text[idx] == "{":
+    start_offset = None
+    end_offset = None
+    lines = text[marker_idx:].splitlines(keepends=True)
+    line_starts = [0]
+    for line in lines[:-1]:
+        line_starts.append(line_starts[-1] + len(line))
+
+    def _offset(row, col):
+        return line_starts[row - 1] + col
+
+    for tok in tokens:
+        if tok.type != tokenize.OP:
+            continue
+        if tok.string == "{":
+            if depth == 0:
+                start_offset = _offset(tok.start[0], tok.start[1])
             depth += 1
-        elif text[idx] == "}":
+        elif tok.string == "}":
             depth -= 1
             if depth == 0:
-                end = idx + 1
+                end_offset = _offset(tok.end[0], tok.end[1])
                 break
-    pyx_targets = ast.literal_eval(text[brace:end])
+    assert start_offset is not None and end_offset is not None, "could not locate SUPPORTED_TARGETS literal"
+    pyx_targets = ast.literal_eval(text[marker_idx + start_offset : marker_idx + end_offset])
 
-    # NVRTC and NVVM map 1:1 to their code_types.
     for backend, code_type in backend_to_code_type.items():
         assert frozenset(pyx_targets[backend]) == _SUPPORTED_TARGETS_BY_CODE_TYPE[code_type], (
             backend,
             code_type,
         )
-    # Both linker backends ultimately serve code_type="ptx"; they should
-    # agree with each other and with the cache table.
     linker_sets = [frozenset(pyx_targets[b]) for b in linker_backends]
     assert all(s == linker_sets[0] for s in linker_sets)
     assert linker_sets[0] == _SUPPORTED_TARGETS_BY_CODE_TYPE["ptx"]