Load pyrefly type info from the types directory

Reviewed By: DinoV

Differential Revision: D95990791

fbshipit-source-id: 4032268c51011f84b4f77525762855e219464c70

Author: martindemello

cinderx/PythonLib/cinderx/compiler/static/pyrefly_compiler.py

@@ -5,20 +5,18 @@
 import sys
 from typing import Callable, Iterable
 
-from cinderx.compiler.static.pyrefly_type_binder import (
-    PyreflyTypeBinder,
-    PyreflyTypeInfo,
-)
+from cinderx.compiler.static.pyrefly_info import EMPTY_TYPE_INFO, Pyrefly
+from cinderx.compiler.static.pyrefly_type_binder import PyreflyTypeBinder
 from cinderx.compiler.static.type_binder import TypeBinder
-from cinderx.compiler.strict.compiler import Compiler, TIMING_LOGGER_TYPE
+from cinderx.compiler.strict.compiler import Compiler
 from cinderx.compiler.strict.flag_extractor import Flags
 from cinderx.compiler.symbols import SymbolVisitor
 
 
 class PyreflyCompiler(Compiler):
     def __init__(
         self,
-        type_info: PyreflyTypeInfo | None = None,
+        pyrefly: Pyrefly | None = None,
         static_opt_out: set[str] | None = None,
         static_opt_in: set[str] | None = None,
         path: Iterable[str] | None = None,
@@ -40,8 +38,7 @@ def __init__(
             use_py_compiler,
             allow_list_regex,
         )
-        assert type_info is not None
-        self.type_info = type_info
+        self.pyrefly = pyrefly
         self.static_opt_in = static_opt_in
         self.static_opt_out = static_opt_out or set()
 
@@ -66,12 +63,15 @@ def make_type_binder(
         optimize: int,
         enable_patching: bool = False,
     ) -> TypeBinder:
+        type_info = EMPTY_TYPE_INFO
+        if self.pyrefly is not None:
+            type_info = self.pyrefly.load_type_info(module_name) or EMPTY_TYPE_INFO
         return PyreflyTypeBinder(
             symbols,
             filename,
             compiler,
             module_name,
             optimize,
             enable_patching,
-            self.type_info,
+            type_info,
         )

cinderx/PythonLib/cinderx/compiler/static/pyrefly_info.py

@@ -0,0 +1,103 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+from __future__ import annotations
+
+import json
+import os
+from ast import AST
+
+
+class PyreflyTypeInfo:
+    """Loads and indexes a pyrefly type trace JSON file.
+
+    The JSON file contains:
+    - type_table: array of type entries (class, literal, callable)
+    - locations: array of {loc: {start_line, start_col, end_line, end_col}, type: index}
+
+    Locations map source positions to type_table indices. Positions use
+    1-based lines and 0-based columns (end_col is exclusive), matching
+    Python AST conventions.
+    """
+
+    def __init__(self, data: dict[str, object]) -> None:
+        self._type_table: list[dict[str, object]] = data["type_table"]
+        self._locations: dict[tuple[int, int, int, int], int] = {}
+        for entry in data["locations"]:
+            loc = entry["loc"]
+            key = (
+                loc["start_line"],
+                loc["start_col"],
+                loc["end_line"],
+                loc["end_col"],
+            )
+            self._locations[key] = entry["type"]
+
+    def _type_to_str(self, type_index: int) -> str:
+        """Convert a type_table entry to a Python annotation string."""
+        entry = self._type_table[type_index]
+        kind = entry["kind"]
+        if kind == "literal":
+            return ""
+        elif kind == "class":
+            qname = str(entry["qname"])
+            args = entry.get("args", [])
+            assert isinstance(args, list)
+            if not args:
+                return qname
+            arg_strs = [self._type_to_str(a) for a in args]
+            if any(not s for s in arg_strs):
+                return qname
+            return f"{qname}[{', '.join(arg_strs)}]"
+        elif kind == "callable":
+            params = entry.get("params", [])
+            assert isinstance(params, list)
+            ret = entry.get("return_type")
+            param_strs = [self._type_to_str(p) for p in params]
+            ret_str = self._type_to_str(ret) if isinstance(ret, int) else ""
+            if any(not s for s in param_strs) or not ret_str:
+                return ""
+            return f"Callable[[{', '.join(param_strs)}], {ret_str}]"
+        return ""
+
+    def lookup(self, node: AST) -> str:
+        """Look up the type string for an AST node by its source position."""
+        key = (
+            node.lineno,  # pyre-ignore[16]
+            node.col_offset,  # pyre-ignore[16]
+            node.end_lineno,  # pyre-ignore[16]
+            node.end_col_offset,  # pyre-ignore[16]
+        )
+        type_index = self._locations.get(key)
+        if type_index is None:
+            return ""
+        return self._type_to_str(type_index)
+
+    @classmethod
+    def load_json(cls, json_path: str) -> PyreflyTypeInfo:
+        with open(json_path) as f:
+            data = json.load(f)
+        return cls(data)
+
+    @classmethod
+    def empty(cls) -> PyreflyTypeInfo:
+        return cls({"type_table": [], "locations": []})
+
+
+class Pyrefly:
+    """Manages type information emitted by pyrefly."""
+
+    def __init__(self, type_dir: str):
+        self.type_dir = type_dir
+
+    def load_type_info(self, module_name: str) -> PyreflyTypeInfo | None:
+        if self.type_dir is None:
+            return None
+        json_path = os.path.join(self.type_dir, "types", f"{module_name}.json")
+        if not os.path.isfile(json_path):
+            return None
+        return PyreflyTypeInfo.load_json(json_path)
+
+
+EMPTY_TYPE_INFO = PyreflyTypeInfo.empty()

cinderx/PythonLib/cinderx/compiler/static/pyrefly_type_binder.py

@@ -12,6 +12,7 @@
 from ..errors import TypedSyntaxError
 from ..symbols import SymbolVisitor
 from .effects import NarrowingEffect
+from .pyrefly_info import PyreflyTypeInfo
 from .type_binder import TypeBinder
 
 if TYPE_CHECKING:

cinderx/PythonLib/cinderx/compiler/strict/pyrefly_loader.py

@@ -14,7 +14,7 @@
 from typing import Callable, Iterable, Mapping
 
 from cinderx.compiler.static.pyrefly_compiler import PyreflyCompiler
-from cinderx.compiler.static.pyrefly_type_binder import PyreflyTypeInfo
+from cinderx.compiler.static.pyrefly_info import Pyrefly
 
 from .compiler import Compiler, TIMING_LOGGER_TYPE
 from .loader import StrictSourceFileLoader
@@ -29,6 +29,8 @@
 
 
 class PyreflyLoader(StrictSourceFileLoader):
+    pyrefly_type_dir: str | None = None
+
     @classmethod
     def ensure_compiler(
         cls,
@@ -41,6 +43,10 @@ def ensure_compiler(
         allow_list_regex: Iterable[str] | None = None,
     ) -> Compiler:
         if (comp := cls.compiler) is None:
+            if cls.pyrefly_type_dir is not None:
+                pyrefly = Pyrefly(cls.pyrefly_type_dir)
+            else:
+                pyrefly = None
             comp = cls.compiler = PyreflyCompiler(
                 type_info=EMPTY_TYPE_INFO,
                 static_opt_out=None,
@@ -49,6 +55,7 @@ def ensure_compiler(
                 stub_path=stub_path,
                 allow_list_prefix=allow_list_prefix,
                 allow_list_exact=allow_list_exact,
+                pyrefly=pyrefly,
                 log_time_func=log_time_func,
                 enable_patching=enable_patching,
                 allow_list_regex=allow_list_regex or [],
@@ -114,8 +121,9 @@ def _get_supported_file_loaders(
     return [extensions, source, bytecode]
 
 
-def install(enable_patching: bool = False) -> None:
+def install(enable_patching: bool = False, pyrefly_type_dir: str | None = None) -> None:
     """Installs a loader which is capable of loading and validating strict modules"""
+    PyreflyLoader.pyrefly_type_dir = pyrefly_type_dir
     supported_loaders = _get_supported_file_loaders(enable_patching)
 
     for index, hook in enumerate(sys.path_hooks):

cinderx/PythonLib/test_cinderx/test_compiler/test_static/pyreflybinder.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import ast
+import unittest
+from types import CodeType
+
+from cinderx.compiler.pycodegen import CodeGenerator
+from cinderx.compiler.static import Compiler, StaticCodeGenBase
+from cinderx.compiler.static.pyrefly_compiler import PyreflyCompiler
+
+from .common import StaticTestBase
+
+
+class PyreBinderTests(StaticTestBase):
+    def test_simple(self):
+        code = self.compile_one("x = 1 + 2")
+
+    def test_force_static(self):
+        compiler = PyreflyCompiler(non_static_modules=set())
+        code, strict, static = compiler.load_compiled_module_from_source(
+            "def f(x: int): return x.is_integer()", "foo.py", "foo", 0
+        )
+        self.assertTrue(static)
+        self.assertFalse(strict)
+        self.assertInBytecode(self.find_code(code, "f"), "INVOKE_METHOD")
+
+    def test_non_force_static(self):
+        compiler = PyreflyCompiler(non_static_modules={"foo"})
+        code, strict, static = compiler.load_compiled_module_from_source(
+            "def f(x: int): return x.is_integer()", "foo.py", "foo", 0
+        )
+        self.assertFalse(static)
+        self.assertFalse(strict)
+        self.assertNotInBytecode(self.find_code(code, "f"), "INVOKE_METHOD")
+
+    def compile_one(
+        self,
+        code: str,
+        modname: str = "<module>",
+        optimize: int = 0,
+        ast_optimizer_enabled: bool = True,
+        enable_patching: bool = False,
+    ) -> CodeType:
+        compiler = PyreflyCompiler()
+        tree = ast.parse(self.clean_code(code))
+        return compiler.compile(
+            modname,
+            f"{modname}.py",
+            tree,
+            code,
+            optimize,
+            enable_patching=enable_patching,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()