Use typed dicts when working with pyrefly json info

Summary: Restores typed dict code after merge issues

Reviewed By: DinoV

Differential Revision: D96209580

fbshipit-source-id: fb6df0661da31ed102b222df99d28ea43ce38fd2

Author: martindemello

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

@@ -7,6 +7,67 @@
 import json
 import os
 from ast import AST
+from dataclasses import dataclass
+from typing import TypedDict
+
+
+class Location(TypedDict):
+    start_line: int
+    start_col: int
+    end_line: int
+    end_col: int
+
+
+TypeKind = int
+
+
+class LocationEntry(TypedDict):
+    loc: LocationInfo
+    type: TypeKind
+
+
+class TypeInfo(TypedDict):
+    type_table: list[TypeTableEntry]
+    locations: list[LocationEntry]
+
+
+class TypeTableEntry(TypedDict):
+    kind: str
+    qname: object
+
+
+@dataclass(slots=True, eq=True)
+class LocationInfo:
+    start_line: int
+    start_col: int
+    end_line: int
+    end_col: int
+
+    def __hash__(self) -> int:
+        return (
+            self.start_col
+            ^ self.start_line << 14
+            ^ self.end_col << 7
+            ^ self.end_line << 32
+        )
+
+    @classmethod
+    def from_location(cls, loc: Location):
+        return cls(
+            start_line=loc["start_line"],
+            start_col=loc["start_col"],
+            end_line=loc["end_line"],
+            end_col=loc["end_col"],
+        )
+
+    @classmethod
+    def from_node(cls, node: AST):
+        return cls(
+            start_line=node.lineno,  # pyre-ignore[16]
+            start_col=node.col_offset,  # pyre-ignore[16]
+            end_line=node.end_lineno,  # pyre-ignore[16]
+            end_col=node.end_col_offset,  # pyre-ignore[16]
+        )
 
 
 class PyreflyTypeInfo:
@@ -21,17 +82,11 @@ class PyreflyTypeInfo:
     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] = {}
+    def __init__(self, data: TypeInfo) -> None:
+        self._type_table: list[TypeTableEntry] = data["type_table"]
+        self._locations: dict[LocationInfo, int] = {}
         for entry in data["locations"]:
-            loc = entry["loc"]
-            key = (
-                loc["start_line"],
-                loc["start_col"],
-                loc["end_line"],
-                loc["end_col"],
-            )
+            key = LocationInfo.from_location(entry["loc"])
             self._locations[key] = entry["type"]
 
     def _type_to_str(self, type_index: int) -> str:
@@ -63,12 +118,7 @@ def _type_to_str(self, type_index: int) -> str:
 
     def _lookup(self, node: AST) -> int | None:
         """Look up the type_table index 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]
-        )
+        key = LocationInfo.from_node(node)
         return self._locations.get(key)
 
     def lookup(self, node: AST) -> str:

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

@@ -6,61 +6,20 @@
 
 import ast
 from ast import AST, Name
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING
 
 from ..errors import TypedSyntaxError
 from ..symbols import SymbolVisitor
 from .effects import NarrowingEffect
 from .module_table import ModuleTable
 from .pyrefly_info import PyreflyTypeInfo
 from .type_binder import TypeBinder
-from .types import Class, Value
+from .types import Class
 
 if TYPE_CHECKING:
     from .compiler import Compiler
 
 
-class Location(TypedDict):
-    start_line: int
-    start_col: int
-    end_line: int
-    end_col: int
-
-
-TypeKind = int
-
-
-class LocationEntry(TypedDict):
-    loc: LocationInfo
-    type: TypeKind
-
-
-class TypeInfo(TypedDict):
-    type_table: list[TypeTable]
-    locations: list[LocationEntry]
-
-
-class TypeTable(TypedDict):
-    kind: str
-    qname: object
-
-
-@dataclass(slots=True, eq=True)
-class LocationInfo:
-    start_line: int
-    start_col: int
-    end_line: int
-    end_col: int
-
-    def __hash__(self) -> int:
-        return (
-            self.start_col
-            ^ self.start_line << 14
-            ^ self.end_col << 7
-            ^ self.end_line << 32
-        )
-
 def _resolve_classname(qname: str, modules: dict[str, ModuleTable]) -> Class | None:
     """Resolve a dotted qname like 'builtins.int' to a Class.
 
@@ -100,72 +59,6 @@ def _resolve_classname(qname: str, modules: dict[str, ModuleTable]) -> Class | N
     return None
 
 
-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, type_info: TypeInfo) -> None:
-        self._type_table: list[TypeTable] = type_info["type_table"]
-        self._locations: dict[LocationInfo, TypeKind] = {}
-        for entry in type_info["locations"]:
-            loc = entry["loc"]
-            key = LocationInfo(
-                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 = LocationInfo(
-            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)
-
-
 class PyreflyTypeBinder(TypeBinder):
     """TypeBinder that uses pyrefly type inference to set types on expression nodes.
 

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

@@ -16,7 +16,7 @@ def test_simple(self):
         code = self.compile_one("x = 1 + 2")
 
     def test_force_static(self):
-        compiler = PyreflyCompiler(non_static_modules=set())
+        compiler = PyreflyCompiler(static_opt_out=set())
         code, strict, static = compiler.load_compiled_module_from_source(
             "def f(x: int): return x.is_integer()", "foo.py", "foo", 0
         )
@@ -25,7 +25,7 @@ def test_force_static(self):
         self.assertInBytecode(self.find_code(code, "f"), "INVOKE_METHOD")
 
     def test_non_force_static(self):
-        compiler = PyreflyCompiler(non_static_modules={"foo"})
+        compiler = PyreflyCompiler(static_opt_out={"foo"})
         code, strict, static = compiler.load_compiled_module_from_source(
             "def f(x: int): return x.is_integer()", "foo.py", "foo", 0
         )