Add Pyrefly loader

Summary: Adds a loader for Pyrefly which uses the PyreflyCompiler

Reviewed By: alexmalyshev

Differential Revision: D96008532

fbshipit-source-id: eb6e079a8edf53c70ce27ce34b767d294cf1c570

Author: DinoV

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

@@ -1,14 +1,16 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
+# pyre-strict
 
 import ast
 import sys
+from typing import Callable, Iterable
 
 from cinderx.compiler.static.pyrefly_type_binder import (
     PyreflyTypeBinder,
     PyreflyTypeInfo,
 )
 from cinderx.compiler.static.type_binder import TypeBinder
-from cinderx.compiler.strict.compiler import Compiler
+from cinderx.compiler.strict.compiler import Compiler, TIMING_LOGGER_TYPE
 from cinderx.compiler.strict.flag_extractor import Flags
 from cinderx.compiler.symbols import SymbolVisitor
 
@@ -17,23 +19,43 @@ class PyreflyCompiler(Compiler):
     def __init__(
         self,
         type_info: PyreflyTypeInfo | None = None,
-        non_static_modules: set[str] | None = None,
+        static_opt_out: set[str] | None = None,
+        static_opt_in: set[str] | None = None,
+        path: Iterable[str] | None = None,
+        stub_path: str | None = None,
+        allow_list_prefix: Iterable[str] | None = None,
+        allow_list_exact: Iterable[str] | None = None,
+        log_time_func: Callable[[], TIMING_LOGGER_TYPE] | None = None,
+        enable_patching: bool = False,
+        use_py_compiler: bool = False,
+        allow_list_regex: Iterable[str] | None = None,
     ):
-        super().__init__(sys.path, "", [], [])
+        super().__init__(
+            path or sys.path,
+            stub_path or "",
+            allow_list_prefix or [],
+            allow_list_exact or [],
+            log_time_func,
+            enable_patching,
+            use_py_compiler,
+            allow_list_regex,
+        )
         assert type_info is not None
         self.type_info = type_info
-        self.non_static_modules = non_static_modules
+        self.static_opt_in = static_opt_in
+        self.static_opt_out = static_opt_out or set()
 
     def get_flags(
         self, module_name: str, pyast: ast.Module, override_flags: Flags
     ) -> Flags:
-        if self.non_static_modules is not None:
-            if module_name in self.non_static_modules:
-                return Flags().merge(override_flags)
+        if self.static_opt_in is not None:
+            return Flags(is_static=module_name in self.static_opt_in).merge(
+                override_flags
+            )
 
-            return Flags(is_static=True).merge(override_flags)
-
-        return super().get_flags(module_name, pyast, override_flags)
+        return Flags(is_static=module_name not in self.static_opt_out).merge(
+            override_flags
+        )
 
     def make_type_binder(
         self,

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

@@ -0,0 +1,130 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# pyre-strict
+
+import sys
+from importlib.abc import Loader
+from importlib.machinery import (
+    BYTECODE_SUFFIXES,
+    EXTENSION_SUFFIXES,
+    ExtensionFileLoader,
+    FileFinder,
+    SOURCE_SUFFIXES,
+    SourcelessFileLoader,
+)
+from typing import Callable, Iterable, Mapping
+
+from cinderx.compiler.static.pyrefly_compiler import PyreflyCompiler
+from cinderx.compiler.static.pyrefly_type_binder import PyreflyTypeInfo
+
+from .compiler import Compiler, TIMING_LOGGER_TYPE
+from .loader import StrictSourceFileLoader
+
+
+EMPTY_TYPE_INFO = PyreflyTypeInfo(
+    {
+        "type_table": [],
+        "locations": [],
+    }
+)
+
+
+class PyreflyLoader(StrictSourceFileLoader):
+    @classmethod
+    def ensure_compiler(
+        cls,
+        path: Iterable[str],
+        stub_path: str,
+        allow_list_prefix: Iterable[str],
+        allow_list_exact: Iterable[str],
+        log_time_func: Callable[[], TIMING_LOGGER_TYPE] | None,
+        enable_patching: bool = False,
+        allow_list_regex: Iterable[str] | None = None,
+    ) -> Compiler:
+        if (comp := cls.compiler) is None:
+            comp = cls.compiler = PyreflyCompiler(
+                type_info=EMPTY_TYPE_INFO,
+                static_opt_out=None,
+                static_opt_in=None,
+                path=path,
+                stub_path=stub_path,
+                allow_list_prefix=allow_list_prefix,
+                allow_list_exact=allow_list_exact,
+                log_time_func=log_time_func,
+                enable_patching=enable_patching,
+                allow_list_regex=allow_list_regex or [],
+            )
+        return comp
+
+    def should_force_strict(self) -> bool:
+        return True
+
+
+class PyreflyLoaderWithPatching(PyreflyLoader):
+    def __init__(
+        self,
+        fullname: str,
+        path: str,
+        import_path: Iterable[str] | None = None,
+        stub_path: str | None = None,
+        allow_list_prefix: Iterable[str] | None = None,
+        allow_list_exact: Iterable[str] | None = None,
+        enable_patching: bool = True,
+        log_source_load: Callable[[str, str | None, bool], None] | None = None,
+        init_cached_properties: None
+        | (
+            Callable[
+                [Mapping[str, str | tuple[str, bool]]],
+                Callable[[type[object]], type[object]],
+            ]
+        ) = None,
+        log_time_func: Callable[[], TIMING_LOGGER_TYPE] | None = None,
+        use_py_compiler: bool = False,
+        # The regexes are parsed on the C++ side, so re.Pattern is not accepted.
+        allow_list_regex: Iterable[str] | None = None,
+    ) -> None:
+        super().__init__(
+            fullname,
+            path,
+            import_path,
+            stub_path,
+            allow_list_prefix,
+            allow_list_exact,
+            enable_patching,
+            log_source_load,
+            init_cached_properties,
+            log_time_func,
+            use_py_compiler,
+            allow_list_regex,
+        )
+
+
+def _get_supported_file_loaders(
+    enable_patching: bool = False,
+) -> list[tuple[type[Loader], list[str]]]:
+    """Returns a list of file-based module loaders.
+
+    Each item is a tuple (loader, suffixes).
+    """
+    extensions = ExtensionFileLoader, EXTENSION_SUFFIXES
+    source = (
+        (PyreflyLoaderWithPatching if enable_patching else PyreflyLoader),
+        SOURCE_SUFFIXES,
+    )
+    bytecode = SourcelessFileLoader, BYTECODE_SUFFIXES
+    return [extensions, source, bytecode]
+
+
+def install(enable_patching: bool = False) -> None:
+    """Installs a loader which is capable of loading and validating strict modules"""
+    supported_loaders = _get_supported_file_loaders(enable_patching)
+
+    for index, hook in enumerate(sys.path_hooks):
+        if not isinstance(hook, type):
+            sys.path_hooks.insert(index, FileFinder.path_hook(*supported_loaders))
+            break
+    else:
+        sys.path_hooks.insert(0, FileFinder.path_hook(*supported_loaders))
+
+    # We need to clear the path_importer_cache so that our new FileFinder will
+    # start being used for existing directories we've loaded modules from.
+    sys.path_importer_cache.clear()

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

@@ -26,7 +26,16 @@ def test_simple(self):
         )
 
     def test_force_static(self):
-        compiler = PyreflyCompiler(non_static_modules=set(), type_info=EMPTY_TYPE_INFO)
+        compiler = PyreflyCompiler(static_opt_out=set(), type_info=EMPTY_TYPE_INFO)
+        code, strict, static = compiler.load_compiled_module_from_source(
+            "def f(x: int): return x.bit_length()", "foo.py", "foo", 0
+        )
+        self.assertTrue(static)
+        self.assertFalse(strict)
+        self.assertIn("<fixed-modules>", code.co_names)
+
+    def test_opt_in_static(self):
+        compiler = PyreflyCompiler(static_opt_in={"foo"}, type_info=EMPTY_TYPE_INFO)
         code, strict, static = compiler.load_compiled_module_from_source(
             "def f(x: int): return x.bit_length()", "foo.py", "foo", 0
         )
@@ -35,9 +44,16 @@ def test_force_static(self):
         self.assertIn("<fixed-modules>", code.co_names)
 
     def test_non_force_static(self):
-        compiler = PyreflyCompiler(
-            non_static_modules={"foo"}, type_info=EMPTY_TYPE_INFO
+        compiler = PyreflyCompiler(static_opt_out={"foo"}, type_info=EMPTY_TYPE_INFO)
+        code, strict, static = compiler.load_compiled_module_from_source(
+            "def f(x: int): return x.bit_length()", "foo.py", "foo", 0
         )
+        self.assertFalse(static)
+        self.assertFalse(strict)
+        self.assertNotIn("<fixed-modules>", code.co_names)
+
+    def test_non_opt_in(self):
+        compiler = PyreflyCompiler(static_opt_in=set(), type_info=EMPTY_TYPE_INFO)
         code, strict, static = compiler.load_compiled_module_from_source(
             "def f(x: int): return x.bit_length()", "foo.py", "foo", 0
         )