Enable native parsing to use source directly

Author: bzoracler

mypy/nativeparse.py

@@ -189,7 +189,10 @@ def add_error(
 
 
 def native_parse(
-    filename: str, options: Options, skip_function_bodies: bool = False
+    filename: str,
+    options: Options,
+    source: str | bytes | None = None,
+    skip_function_bodies: bool = False,
 ) -> tuple[MypyFile, list[ParseError], TypeIgnores]:
     """Parse a Python file using the native Rust-based parser.
 
@@ -218,7 +221,7 @@ def native_parse(
         uses_template_strings,
         source_hash,
         mypy_comments,
-    ) = parse_to_binary_ast(filename, options, skip_function_bodies)
+    ) = parse_to_binary_ast(filename, options, source, skip_function_bodies)
     node = MypyFile([], [])
     node.path = filename
     node.raw_data = FileRawData(
@@ -255,7 +258,10 @@ def read_statements(state: State, data: ReadBuffer, n: int) -> list[Statement]:
 
 
 def parse_to_binary_ast(
-    filename: str, options: Options, skip_function_bodies: bool = False
+    filename: str,
+    options: Options,
+    source: str | bytes | None = None,
+    skip_function_bodies: bool = False,
 ) -> tuple[bytes, list[ParseError], TypeIgnores, bytes, bool, bool, str, list[tuple[int, str]]]:
     # This is a horrible hack to work around a mypyc bug where imported
     # module may be not ready in a thread sometimes.
@@ -266,6 +272,7 @@ def parse_to_binary_ast(
             raise ImportError("Cannot import ast_serialize")
     ast_bytes, errors, ignores, import_bytes, ast_data = ast_serialize.parse(
         filename,
+        source,
         skip_function_bodies=skip_function_bodies,
         python_version=options.python_version,
         platform=options.platform,

mypy/test/test_nativeparse.py

@@ -234,7 +234,7 @@ def format_reachable_imports(node: MypyFile) -> list[str]:
 
 @unittest.skipUnless(has_nativeparse, "nativeparse not available")
 class TestNativeParserBinaryFormat(unittest.TestCase):
-    def test_trivial_binary_data(self) -> None:
+    def _assert_trivial_binary_data(self, b: bytes, /) -> None:
         # A quick sanity check to ensure the serialized data looks as expected. Only covers
         # a few AST nodes.
 
@@ -250,9 +250,9 @@ def locs(start_line: int, start_column: int, end_line: int, end_column: int) ->
                 int_enc(end_column - start_column),
             ]
 
-        with temp_source("print('hello')") as fnam:
-            b, _, _, _, _, _, _, _ = parse_to_binary_ast(fnam, Options())
-            assert list(b) == (
+        self.assertEqual(
+            list(b),
+            (
                 [LITERAL_INT, 22, nodes.EXPR_STMT, nodes.CALL_EXPR]
                 + [nodes.NAME_EXPR, LITERAL_STR]
                 + [int_enc(5)]
@@ -269,7 +269,21 @@ def locs(start_line: int, start_column: int, end_line: int, end_column: int) ->
                 + [LIST_GEN, 22, LITERAL_NONE]
                 + locs(1, 0, 1, 14)
                 + [END_TAG, END_TAG]
-            )
+            ),
+        )
+
+    def test_trivial_binary_data_from_file(self) -> None:
+        with temp_source("print('hello')") as fnam:
+            b, _, _, _, _, _, _, _ = parse_to_binary_ast(fnam, Options())
+            self._assert_trivial_binary_data(b)
+
+    def test_trivial_binary_data_from_string_source(self) -> None:
+        b, _, _, _, _, _, _, _ = parse_to_binary_ast("", Options(), "print('hello')")
+        self._assert_trivial_binary_data(b)
+
+    def test_trivial_binary_data_from_bytes_source(self) -> None:
+        b, _, _, _, _, _, _, _ = parse_to_binary_ast("", Options(), b"print('hello')")
+        self._assert_trivial_binary_data(b)
 
 
 @contextlib.contextmanager