Python: Output JSON from tsg-python

Changes the current text-based dump of the tsg-python graph to the
native JSON output that tree-sitter-graph provides. This saves us the
hassle of parsing the output ourselves, including the recently added bug
fix that allowed us to interpret Rust-style string escapes in Python.

As an added bonus, relying on the json module instead of our own parser
roughly cuts the AST reconstruction time in half compared to previously,
an extraction speedup in the neighbourhood of 3-4% (when tested on an
extraction of `python/cpython` using only the tsg-python parser).

Author: tausbn

MODULE.bazel

@@ -93,6 +93,7 @@ use_repo(
     "vendor_py__cc-1.2.14",
     "vendor_py__clap-4.5.30",
     "vendor_py__regex-1.11.1",
+    "vendor_py__serde_json-1.0.138",
     "vendor_py__tree-sitter-0.24.7",
     "vendor_py__tree-sitter-graph-0.12.0",
 )

misc/bazel/3rdparty/py_deps/BUILD.bazel

@@ -3,7 +3,7 @@
 # Functions and classes used for parsing Python files using `tree-sitter-graph`
 
 from ast import literal_eval
-import re
+import json
 import sys
 import os
 import semmle.python.parser
@@ -114,98 +114,67 @@ def __repr__(self):
     cargo_file = os.path.join(tsg_python_path, "Cargo.toml")
     tsg_command = ["cargo", "run", "--quiet", "--release", "--manifest-path="+cargo_file]
 
+def _decode_tsg_value(value):
+    value_type = value["type"]
+    if value_type == "null":
+        return None
+    if value_type in ("bool", "int", "string"):
+        return value[value_type]
+    if value_type == "list":
+        return [_decode_tsg_value(item) for item in value["values"]]
+    if value_type == "graphNode":
+        return Node(value["id"])
+    raise ValueError("Unsupported TSG value type '{}'".format(value_type))
+
+def _decode_tsg_node_attributes(encoded_attrs, path, logger):
+    # Decode every attribute first so location information is available if string decoding fails.
+    attrs = {
+        key: _decode_tsg_value(value)
+        for key, value in encoded_attrs.items()
+    }
+    if "s" not in attrs:
+        return attrs
+
+    try:
+        attrs["s"] = evaluate_string(attrs["s"])
+    except Exception as ex:
+        loc = ":".join(str(i) for i in get_location_info(attrs))
+        error = ex.args[0] if ex.args else "unknown"
+        logger.warning(
+            "Error '{}' while parsing value {} at {}:{}\n".format(
+                error, repr(attrs["s"]), path, loc
+            )
+        )
+    return attrs
+
 def read_tsg_python_output(path, logger):
+    command_args = tsg_command + [path]
+    p = subprocess.Popen(command_args, stdout=subprocess.PIPE)
+    stdout, _ = p.communicate()
+    if p.returncode:
+        raise subprocess.CalledProcessError(p.returncode, command_args, stdout)
+
     # Mapping from node id (an integer) to a dictionary containing attribute data.
     node_attr = {}
     # Mapping a start node to a map from attribute names to lists of (value, end_node) pairs.
     edge_attr = {}
 
-    command_args = tsg_command + [path]
-    p = subprocess.Popen(command_args, stdout=subprocess.PIPE)
-    for line in p.stdout:
-        line = line.decode(sys.getfilesystemencoding())
-        line = line.rstrip()
-        if line.startswith("node"): # e.g. `node 5`
-            current_node = int(line.split(" ")[1])
-            d = {}
-            node_attr[current_node] = d
-            in_node = True
-        elif line.startswith("edge"): # e.g. `edge 5 -> 6`
-            current_start, current_end = tuple(map(int, line[4:].split("->")))
-            d = edge_attr.setdefault(current_start, {})
-            in_node = False
-        else: # attribute, e.g. `_kind: "Class"`
-            key, value = line[2:].split(": ", 1)
-            if value.startswith("[graph node"): # e.g. `_skip_to: [graph node 5]`
-                value = Node(int(value.split(" ")[2][:-1]))
-            elif value == "#true": # e.g. `_is_parenthesised: #true`
-                value = True
-            elif value == "#false": # e.g. `top: #false`
-                value = False
-            elif value == "#null": # e.g. `exc: #null`
-                value = None
-            else: # literal values, e.g. `name: "k1.k2"` or `level: 5`
-                value = rust_to_python_escapes(value)
-                try:
-                    if key =="s" and value[0] == '"': # e.g. `s: "k1.k2"`
-                        value = evaluate_string(value)
-                    else:
-                        value  = literal_eval(value)
-                        if isinstance(value, bytes):
-                            try:
-                                value = value.decode(sys.getfilesystemencoding())
-                            except UnicodeDecodeError:
-                                # just include the bytes as-is
-                                pass
-                except Exception as ex:
-                    # We may not know the location at this point -- for instance if we forgot to set
-                    # it -- but `get_location_info` will degrade gracefully in this case.
-                    loc = ":".join(str(i) for i in get_location_info(d))
-                    error = ex.args[0] if ex.args else "unknown"
-                    logger.warning("Error '{}' while parsing value {} at {}:{}\n".format(error, repr(value), path, loc))
-            if in_node:
-                d[key] = value
-            else:
-                d.setdefault(key, []).append((value, current_end))
-    p.stdout.close()
-    p.terminate()
-    p.wait()
+    for encoded_node in json.loads(stdout):
+        current_node = encoded_node["id"]
+        attrs = _decode_tsg_node_attributes(encoded_node["attrs"], path, logger)
+        node_attr[current_node] = attrs
+        for encoded_edge in encoded_node["edges"]:
+            current_end = encoded_edge["sink"]
+            edge_fields = edge_attr.setdefault(current_node, {})
+            for key, value in encoded_edge["attrs"].items():
+                value = _decode_tsg_value(value)
+                edge_fields.setdefault(key, []).append((value, current_end))
     logger.debug("Read {} nodes and {} edges from TSG output".format(len(node_attr), len(edge_attr)))
     return node_attr, edge_attr
 
-# `tsg-python` serialises string values using Rust's `Debug` formatting, which diverges from what
-# Python's `literal_eval` accepts in two ways:
-#  - characters Rust considers non-printable -- including grapheme-extending ones such as the U+FE0F
-#    variation selector, U+200D zero width joiner and combining accents -- are rendered as `\u{...}`,
-#    a syntax Python does not know at all;
-#  - NUL is rendered as `\0`, which Python reads as the start of an *octal* escape, silently
-#    swallowing up to two more digits (NUL followed by `1` is emitted as `"\01"`, which decodes
-#    to `\x01`).
-# Everything else Rust emits (`\t`, `\r`, `\n`, `\\`, `\"`, and unescaped characters) is read back
-# identically by `literal_eval`, as verified exhaustively over every Unicode scalar value.
-_RUST_ESCAPE = re.compile(r"\\(?:u\{([0-9a-fA-F]{1,6})\}|.)", re.DOTALL)
-
-def rust_to_python_escapes(text):
-    """Rewrites Rust escapes in `text` that Python would reject or misread into their equivalents.
-
-    Matching every escape sequence (rather than only the offending ones) keeps the scan in step with
-    the backslashes, so an escaped backslash -- how a literal `\\u{fe0f}` in the source is
-    serialised -- is left alone."""
-    if "\\u{" not in text and "\\0" not in text:
-        return text
-    def replace(match):
-        code_point = match.group(1)
-        if code_point is None:
-            return "\\x00" if match.group(0) == "\\0" else match.group(0)
-        code_point = int(code_point, 16)
-        if code_point > 0xFFFF:
-            return "\\U{:08x}".format(code_point)
-        return "\\u{:04x}".format(code_point)
-    return _RUST_ESCAPE.sub(replace, text)
-
-def evaluate_string(s):
-    s = literal_eval(s)
-    prefix, quotes, content = split_string(s, None)
+def evaluate_string(source_literal):
+    """Evaluates Python string literal text that has already been decoded from the wire format."""
+    prefix, quotes, content = split_string(source_literal, None)
     ends_with_illegal_character = False
     # If the string ends with the same quote character as the outer quotes (and/or backslashes)
     # (e.g. the first string part of `f"""hello"{0}"""`), we must take care to not accidently create

misc/bazel/3rdparty/py_deps/defs.bzl

@@ -10,7 +10,7 @@
 
 #Semantic version of extractor.
 #Update this if any changes are made
-VERSION = "7.1.10"
+VERSION = "7.1.11"
 
 PY_EXTENSIONS = ".py", ".pyw"
 

python/extractor/semmle/python/parser/tsg_parser.py

@@ -1,75 +1,76 @@
 import unittest
-
-from ast import literal_eval
+import json
 
 from semmle.logging import format_message
-from semmle.python.parser.tsg_parser import evaluate_string, rust_to_python_escapes
-
-
-class RustEscapeTest(unittest.TestCase):
-    """`tsg-python` serialises strings with Rust's `Debug` formatting, which escapes characters such
-    as U+FE0F as `\\u{...}` -- a syntax Python's `literal_eval` does not accept -- and NUL as `\\0`,
-    which Python reads as an octal escape."""
-
-    def test_untouched_without_escapes(self):
-        text = '"caf\u00e9 \u2713 \U0001f4be"'
-        self.assertEqual(rust_to_python_escapes(text), text)
-
-    def test_basic_multilingual_plane(self):
-        self.assertEqual(rust_to_python_escapes(r'"\u{fe0f}"'), r'"\ufe0f"')
-        self.assertEqual(rust_to_python_escapes(r'"\u{200d}"'), r'"\u200d"')
-
-    def test_short_and_astral_code_points(self):
-        self.assertEqual(rust_to_python_escapes(r'"\u{0}"'), r'"\u0000"')
-        self.assertEqual(rust_to_python_escapes(r'"\u{1f4a9}"'), r'"\U0001f4a9"')
-
-    def test_other_escapes_are_preserved(self):
-        self.assertEqual(rust_to_python_escapes(r'"a\nb\"c\u{ad}"'), r'"a\nb\"c\u00ad"')
-
-    def test_escaped_backslash_is_not_an_escape_introducer(self):
-        # How a raw string `r"\u{fe0f}"` in the analysed source gets serialised: the `\u{fe0f}` is
-        # literal text, not an escape, and must survive unchanged.
-        self.assertEqual(rust_to_python_escapes(r'"\\u{fe0f}"'), r'"\\u{fe0f}"')
-
-    def test_nul_is_not_left_as_an_octal_escape(self):
-        # Rust renders NUL as `\0`; Python would read that as the start of an octal escape and
-        # swallow the digits that follow, decoding `"\01"` to U+0001 instead of NUL then `1`.
-        self.assertEqual(rust_to_python_escapes(r'"\01"'), r'"\x001"')
-
-    def test_every_escape_shape_round_trips(self):
-        # Rust's `Debug for str` only ever emits these escape shapes. Check that each round-trips
-        # with every printable ASCII neighbour before and after it.
-        for escape_shape, expected in [
-            (r'\0', "\x00"),
-            (r'\t', "\t"),
-            (r'\n', "\n"),
-            (r'\r', "\r"),
-            (r'\\', "\\"),
-            (r'\"', '"'),
-            (r'\u{1}', "\u0001"),
-            (r'\u{1f}', "\u001f"),
-            (r'\u{300}', "\u0300"),
-            (r'\u{fe0f}', "\ufe0f"),
-            (r'\u{e0100}', "\U000e0100"),
-            (r'\u{10fffe}', "\U0010fffe"),
-        ]:
-            for neighbour in map(chr, range(0x20, 0x7F)):
-                rendered_neighbour = {"\\": r"\\", '"': r'\"'}.get(neighbour, neighbour)
-                for position, text, expected_value in [
-                    ("before", '"' + rendered_neighbour + escape_shape + '"', neighbour + expected),
-                    ("after", '"' + escape_shape + rendered_neighbour + '"', expected + neighbour),
-                ]:
-                    with self.subTest(
-                        escape_shape=escape_shape,
-                        neighbour=neighbour,
-                        position=position,
-                    ):
-                        self.assertEqual(literal_eval(rust_to_python_escapes(text)), expected_value)
+from semmle.python.parser.tsg_parser import Node, evaluate_string, read_tsg_python_output
+
+
+class JsonOutputTest(unittest.TestCase):
+    def test_decodes_nodes_edges_and_attribute_values(self):
+        output = json.dumps(
+            [
+                {
+                    "id": 0,
+                    "edges": [
+                        {
+                            "sink": 1,
+                            "attrs": {"body": {"type": "int", "int": 0}},
+                        }
+                    ],
+                    "attrs": {
+                        "_kind": {"type": "string", "string": "Module"},
+                        "_location": {
+                            "type": "list",
+                            "values": [
+                                {"type": "int", "int": 0},
+                                {"type": "int", "int": 0},
+                                {"type": "int", "int": 1},
+                                {"type": "int", "int": 0},
+                            ],
+                        },
+                    },
+                },
+                {
+                    "id": 1,
+                    "edges": [],
+                    "attrs": {
+                        "_kind": {"type": "string", "string": "Name"},
+                        "variable": {
+                            "type": "string",
+                            "string": "caf\u00e9 \u26a0\ufe0f \U0001f4be",
+                        },
+                        "s": {
+                            "type": "string",
+                            "string": '"\u26a0\ufe0f  problem %s: %s"',
+                        },
+                        "is_async": {"type": "bool", "bool": True},
+                        "optional": {"type": "null"},
+                        "_skip_to": {"type": "graphNode", "id": 0},
+                    },
+                },
+            ]
+        ).encode("utf-8")
+
+        process = unittest.mock.Mock()
+        process.communicate.return_value = (output, None)
+        process.returncode = 0
+        with unittest.mock.patch(
+            "semmle.python.parser.tsg_parser.subprocess.Popen", return_value=process
+        ):
+            node_attr, edge_attr = read_tsg_python_output(
+                "test.py", unittest.mock.Mock()
+            )
+
+        self.assertEqual(node_attr[1]["variable"], "caf\u00e9 \u26a0\ufe0f \U0001f4be")
+        self.assertEqual(node_attr[1]["s"], "\u26a0\ufe0f  problem %s: %s")
+        self.assertIs(node_attr[1]["is_async"], True)
+        self.assertIsNone(node_attr[1]["optional"])
+        self.assertIsInstance(node_attr[1]["_skip_to"], Node)
+        self.assertEqual(node_attr[1]["_skip_to"].id, 0)
+        self.assertEqual(edge_attr, {0: {"body": [(0, 1)]}})
 
     def test_evaluate_string_on_reported_value(self):
-        # The exact value from https://github.com/github/codeql/issues/22435 that used to raise
-        # `truncated \uXXXX escape`.
-        value = rust_to_python_escapes('"\\"\u26a0\\u{fe0f}  problem %s: %s\\""')
+        value = '"\u26a0\ufe0f  problem %s: %s"'
         self.assertEqual(evaluate_string(value), "\u26a0\ufe0f  problem %s: %s")
 
 

python/extractor/semmle/util.py

@@ -14,3 +14,4 @@ tree-sitter = "=0.24.7"
 tree-sitter-graph = "0.12.0"
 tsp = {path = "tsp"}
 clap = "4.5"
+serde_json = "1.0"

python/extractor/tests/test_tsg_parser.py

@@ -710,9 +710,6 @@ fn main() -> Result<()> {
         add_syntax_error_nodes(&mut graph, &syntax_errors);
     }
 
-    // `pretty_print` renders string values with Rust's `Debug` formatting, so non-printable and
-    // grapheme-extending characters come out as `\u{...}`. The reader on the other side
-    // (`semmle/python/parser/tsg_parser.py`) translates those into Python escapes.
-    print!("{}", graph.pretty_print());
+    serde_json::to_writer(std::io::stdout().lock(), &graph)?;
     Ok(())
 }