Skip to content

Commit aa5e474

Browse files
committed
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).
1 parent 53db3bd commit aa5e474

9 files changed

Lines changed: 140 additions & 156 deletions

File tree

MODULE.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ use_repo(
9393
"vendor_py__cc-1.2.14",
9494
"vendor_py__clap-4.5.30",
9595
"vendor_py__regex-1.11.1",
96+
"vendor_py__serde_json-1.0.138",
9697
"vendor_py__tree-sitter-0.24.7",
9798
"vendor_py__tree-sitter-graph-0.12.0",
9899
)

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

Lines changed: 12 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 53 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
# Functions and classes used for parsing Python files using `tree-sitter-graph`
44

55
from ast import literal_eval
6-
import re
6+
import json
77
import sys
88
import os
99
import semmle.python.parser
@@ -114,98 +114,67 @@ def __repr__(self):
114114
cargo_file = os.path.join(tsg_python_path, "Cargo.toml")
115115
tsg_command = ["cargo", "run", "--quiet", "--release", "--manifest-path="+cargo_file]
116116

117+
def _decode_tsg_value(value):
118+
value_type = value["type"]
119+
if value_type == "null":
120+
return None
121+
if value_type in ("bool", "int", "string"):
122+
return value[value_type]
123+
if value_type == "list":
124+
return [_decode_tsg_value(item) for item in value["values"]]
125+
if value_type == "graphNode":
126+
return Node(value["id"])
127+
raise ValueError("Unsupported TSG value type '{}'".format(value_type))
128+
129+
def _decode_tsg_node_attributes(encoded_attrs, path, logger):
130+
# Decode every attribute first so location information is available if string decoding fails.
131+
attrs = {
132+
key: _decode_tsg_value(value)
133+
for key, value in encoded_attrs.items()
134+
}
135+
if "s" not in attrs:
136+
return attrs
137+
138+
try:
139+
attrs["s"] = evaluate_string(attrs["s"])
140+
except Exception as ex:
141+
loc = ":".join(str(i) for i in get_location_info(attrs))
142+
error = ex.args[0] if ex.args else "unknown"
143+
logger.warning(
144+
"Error '{}' while parsing value {} at {}:{}\n".format(
145+
error, repr(attrs["s"]), path, loc
146+
)
147+
)
148+
return attrs
149+
117150
def read_tsg_python_output(path, logger):
151+
command_args = tsg_command + [path]
152+
p = subprocess.Popen(command_args, stdout=subprocess.PIPE)
153+
stdout, _ = p.communicate()
154+
if p.returncode:
155+
raise subprocess.CalledProcessError(p.returncode, command_args, stdout)
156+
118157
# Mapping from node id (an integer) to a dictionary containing attribute data.
119158
node_attr = {}
120159
# Mapping a start node to a map from attribute names to lists of (value, end_node) pairs.
121160
edge_attr = {}
122161

123-
command_args = tsg_command + [path]
124-
p = subprocess.Popen(command_args, stdout=subprocess.PIPE)
125-
for line in p.stdout:
126-
line = line.decode(sys.getfilesystemencoding())
127-
line = line.rstrip()
128-
if line.startswith("node"): # e.g. `node 5`
129-
current_node = int(line.split(" ")[1])
130-
d = {}
131-
node_attr[current_node] = d
132-
in_node = True
133-
elif line.startswith("edge"): # e.g. `edge 5 -> 6`
134-
current_start, current_end = tuple(map(int, line[4:].split("->")))
135-
d = edge_attr.setdefault(current_start, {})
136-
in_node = False
137-
else: # attribute, e.g. `_kind: "Class"`
138-
key, value = line[2:].split(": ", 1)
139-
if value.startswith("[graph node"): # e.g. `_skip_to: [graph node 5]`
140-
value = Node(int(value.split(" ")[2][:-1]))
141-
elif value == "#true": # e.g. `_is_parenthesised: #true`
142-
value = True
143-
elif value == "#false": # e.g. `top: #false`
144-
value = False
145-
elif value == "#null": # e.g. `exc: #null`
146-
value = None
147-
else: # literal values, e.g. `name: "k1.k2"` or `level: 5`
148-
value = rust_to_python_escapes(value)
149-
try:
150-
if key =="s" and value[0] == '"': # e.g. `s: "k1.k2"`
151-
value = evaluate_string(value)
152-
else:
153-
value = literal_eval(value)
154-
if isinstance(value, bytes):
155-
try:
156-
value = value.decode(sys.getfilesystemencoding())
157-
except UnicodeDecodeError:
158-
# just include the bytes as-is
159-
pass
160-
except Exception as ex:
161-
# We may not know the location at this point -- for instance if we forgot to set
162-
# it -- but `get_location_info` will degrade gracefully in this case.
163-
loc = ":".join(str(i) for i in get_location_info(d))
164-
error = ex.args[0] if ex.args else "unknown"
165-
logger.warning("Error '{}' while parsing value {} at {}:{}\n".format(error, repr(value), path, loc))
166-
if in_node:
167-
d[key] = value
168-
else:
169-
d.setdefault(key, []).append((value, current_end))
170-
p.stdout.close()
171-
p.terminate()
172-
p.wait()
162+
for encoded_node in json.loads(stdout):
163+
current_node = encoded_node["id"]
164+
attrs = _decode_tsg_node_attributes(encoded_node["attrs"], path, logger)
165+
node_attr[current_node] = attrs
166+
for encoded_edge in encoded_node["edges"]:
167+
current_end = encoded_edge["sink"]
168+
edge_fields = edge_attr.setdefault(current_node, {})
169+
for key, value in encoded_edge["attrs"].items():
170+
value = _decode_tsg_value(value)
171+
edge_fields.setdefault(key, []).append((value, current_end))
173172
logger.debug("Read {} nodes and {} edges from TSG output".format(len(node_attr), len(edge_attr)))
174173
return node_attr, edge_attr
175174

176-
# `tsg-python` serialises string values using Rust's `Debug` formatting, which diverges from what
177-
# Python's `literal_eval` accepts in two ways:
178-
# - characters Rust considers non-printable -- including grapheme-extending ones such as the U+FE0F
179-
# variation selector, U+200D zero width joiner and combining accents -- are rendered as `\u{...}`,
180-
# a syntax Python does not know at all;
181-
# - NUL is rendered as `\0`, which Python reads as the start of an *octal* escape, silently
182-
# swallowing up to two more digits (NUL followed by `1` is emitted as `"\01"`, which decodes
183-
# to `\x01`).
184-
# Everything else Rust emits (`\t`, `\r`, `\n`, `\\`, `\"`, and unescaped characters) is read back
185-
# identically by `literal_eval`, as verified exhaustively over every Unicode scalar value.
186-
_RUST_ESCAPE = re.compile(r"\\(?:u\{([0-9a-fA-F]{1,6})\}|.)", re.DOTALL)
187-
188-
def rust_to_python_escapes(text):
189-
"""Rewrites Rust escapes in `text` that Python would reject or misread into their equivalents.
190-
191-
Matching every escape sequence (rather than only the offending ones) keeps the scan in step with
192-
the backslashes, so an escaped backslash -- how a literal `\\u{fe0f}` in the source is
193-
serialised -- is left alone."""
194-
if "\\u{" not in text and "\\0" not in text:
195-
return text
196-
def replace(match):
197-
code_point = match.group(1)
198-
if code_point is None:
199-
return "\\x00" if match.group(0) == "\\0" else match.group(0)
200-
code_point = int(code_point, 16)
201-
if code_point > 0xFFFF:
202-
return "\\U{:08x}".format(code_point)
203-
return "\\u{:04x}".format(code_point)
204-
return _RUST_ESCAPE.sub(replace, text)
205-
206-
def evaluate_string(s):
207-
s = literal_eval(s)
208-
prefix, quotes, content = split_string(s, None)
175+
def evaluate_string(source_literal):
176+
"""Evaluates Python string literal text that has already been decoded from the wire format."""
177+
prefix, quotes, content = split_string(source_literal, None)
209178
ends_with_illegal_character = False
210179
# If the string ends with the same quote character as the outer quotes (and/or backslashes)
211180
# (e.g. the first string part of `f"""hello"{0}"""`), we must take care to not accidently create

python/extractor/semmle/util.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
#Semantic version of extractor.
1212
#Update this if any changes are made
13-
VERSION = "7.1.10"
13+
VERSION = "7.1.11"
1414

1515
PY_EXTENSIONS = ".py", ".pyw"
1616

python/extractor/tests/test_tsg_parser.py

Lines changed: 68 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -1,75 +1,76 @@
11
import unittest
2-
3-
from ast import literal_eval
2+
import json
43

54
from semmle.logging import format_message
6-
from semmle.python.parser.tsg_parser import evaluate_string, rust_to_python_escapes
7-
8-
9-
class RustEscapeTest(unittest.TestCase):
10-
"""`tsg-python` serialises strings with Rust's `Debug` formatting, which escapes characters such
11-
as U+FE0F as `\\u{...}` -- a syntax Python's `literal_eval` does not accept -- and NUL as `\\0`,
12-
which Python reads as an octal escape."""
13-
14-
def test_untouched_without_escapes(self):
15-
text = '"caf\u00e9 \u2713 \U0001f4be"'
16-
self.assertEqual(rust_to_python_escapes(text), text)
17-
18-
def test_basic_multilingual_plane(self):
19-
self.assertEqual(rust_to_python_escapes(r'"\u{fe0f}"'), r'"\ufe0f"')
20-
self.assertEqual(rust_to_python_escapes(r'"\u{200d}"'), r'"\u200d"')
21-
22-
def test_short_and_astral_code_points(self):
23-
self.assertEqual(rust_to_python_escapes(r'"\u{0}"'), r'"\u0000"')
24-
self.assertEqual(rust_to_python_escapes(r'"\u{1f4a9}"'), r'"\U0001f4a9"')
25-
26-
def test_other_escapes_are_preserved(self):
27-
self.assertEqual(rust_to_python_escapes(r'"a\nb\"c\u{ad}"'), r'"a\nb\"c\u00ad"')
28-
29-
def test_escaped_backslash_is_not_an_escape_introducer(self):
30-
# How a raw string `r"\u{fe0f}"` in the analysed source gets serialised: the `\u{fe0f}` is
31-
# literal text, not an escape, and must survive unchanged.
32-
self.assertEqual(rust_to_python_escapes(r'"\\u{fe0f}"'), r'"\\u{fe0f}"')
33-
34-
def test_nul_is_not_left_as_an_octal_escape(self):
35-
# Rust renders NUL as `\0`; Python would read that as the start of an octal escape and
36-
# swallow the digits that follow, decoding `"\01"` to U+0001 instead of NUL then `1`.
37-
self.assertEqual(rust_to_python_escapes(r'"\01"'), r'"\x001"')
38-
39-
def test_every_escape_shape_round_trips(self):
40-
# Rust's `Debug for str` only ever emits these escape shapes. Check that each round-trips
41-
# with every printable ASCII neighbour before and after it.
42-
for escape_shape, expected in [
43-
(r'\0', "\x00"),
44-
(r'\t', "\t"),
45-
(r'\n', "\n"),
46-
(r'\r', "\r"),
47-
(r'\\', "\\"),
48-
(r'\"', '"'),
49-
(r'\u{1}', "\u0001"),
50-
(r'\u{1f}', "\u001f"),
51-
(r'\u{300}', "\u0300"),
52-
(r'\u{fe0f}', "\ufe0f"),
53-
(r'\u{e0100}', "\U000e0100"),
54-
(r'\u{10fffe}', "\U0010fffe"),
55-
]:
56-
for neighbour in map(chr, range(0x20, 0x7F)):
57-
rendered_neighbour = {"\\": r"\\", '"': r'\"'}.get(neighbour, neighbour)
58-
for position, text, expected_value in [
59-
("before", '"' + rendered_neighbour + escape_shape + '"', neighbour + expected),
60-
("after", '"' + escape_shape + rendered_neighbour + '"', expected + neighbour),
61-
]:
62-
with self.subTest(
63-
escape_shape=escape_shape,
64-
neighbour=neighbour,
65-
position=position,
66-
):
67-
self.assertEqual(literal_eval(rust_to_python_escapes(text)), expected_value)
5+
from semmle.python.parser.tsg_parser import Node, evaluate_string, read_tsg_python_output
6+
7+
8+
class JsonOutputTest(unittest.TestCase):
9+
def test_decodes_nodes_edges_and_attribute_values(self):
10+
output = json.dumps(
11+
[
12+
{
13+
"id": 0,
14+
"edges": [
15+
{
16+
"sink": 1,
17+
"attrs": {"body": {"type": "int", "int": 0}},
18+
}
19+
],
20+
"attrs": {
21+
"_kind": {"type": "string", "string": "Module"},
22+
"_location": {
23+
"type": "list",
24+
"values": [
25+
{"type": "int", "int": 0},
26+
{"type": "int", "int": 0},
27+
{"type": "int", "int": 1},
28+
{"type": "int", "int": 0},
29+
],
30+
},
31+
},
32+
},
33+
{
34+
"id": 1,
35+
"edges": [],
36+
"attrs": {
37+
"_kind": {"type": "string", "string": "Name"},
38+
"variable": {
39+
"type": "string",
40+
"string": "caf\u00e9 \u26a0\ufe0f \U0001f4be",
41+
},
42+
"s": {
43+
"type": "string",
44+
"string": '"\u26a0\ufe0f problem %s: %s"',
45+
},
46+
"is_async": {"type": "bool", "bool": True},
47+
"optional": {"type": "null"},
48+
"_skip_to": {"type": "graphNode", "id": 0},
49+
},
50+
},
51+
]
52+
).encode("utf-8")
53+
54+
process = unittest.mock.Mock()
55+
process.communicate.return_value = (output, None)
56+
process.returncode = 0
57+
with unittest.mock.patch(
58+
"semmle.python.parser.tsg_parser.subprocess.Popen", return_value=process
59+
):
60+
node_attr, edge_attr = read_tsg_python_output(
61+
"test.py", unittest.mock.Mock()
62+
)
63+
64+
self.assertEqual(node_attr[1]["variable"], "caf\u00e9 \u26a0\ufe0f \U0001f4be")
65+
self.assertEqual(node_attr[1]["s"], "\u26a0\ufe0f problem %s: %s")
66+
self.assertIs(node_attr[1]["is_async"], True)
67+
self.assertIsNone(node_attr[1]["optional"])
68+
self.assertIsInstance(node_attr[1]["_skip_to"], Node)
69+
self.assertEqual(node_attr[1]["_skip_to"].id, 0)
70+
self.assertEqual(edge_attr, {0: {"body": [(0, 1)]}})
6871

6972
def test_evaluate_string_on_reported_value(self):
70-
# The exact value from https://github.com/github/codeql/issues/22435 that used to raise
71-
# `truncated \uXXXX escape`.
72-
value = rust_to_python_escapes('"\\"\u26a0\\u{fe0f} problem %s: %s\\""')
73+
value = '"\u26a0\ufe0f problem %s: %s"'
7374
self.assertEqual(evaluate_string(value), "\u26a0\ufe0f problem %s: %s")
7475

7576

python/extractor/tsg-python/Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

python/extractor/tsg-python/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,4 @@ tree-sitter = "=0.24.7"
1414
tree-sitter-graph = "0.12.0"
1515
tsp = {path = "tsp"}
1616
clap = "4.5"
17+
serde_json = "1.0"

python/extractor/tsg-python/src/main.rs

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -710,9 +710,6 @@ fn main() -> Result<()> {
710710
add_syntax_error_nodes(&mut graph, &syntax_errors);
711711
}
712712

713-
// `pretty_print` renders string values with Rust's `Debug` formatting, so non-printable and
714-
// grapheme-extending characters come out as `\u{...}`. The reader on the other side
715-
// (`semmle/python/parser/tsg_parser.py`) translates those into Python escapes.
716-
print!("{}", graph.pretty_print());
713+
serde_json::to_writer(std::io::stdout().lock(), &graph)?;
717714
Ok(())
718715
}

0 commit comments

Comments
 (0)