More Pyrefly binding support

Summary: The static compiler has some additional metadata beyond types that it needs. This fills in those details so that we're able to get _pytree compiling. Also includes a fix for protocol methods so that we don't complain that they have the wrong return type.

Reviewed By: martindemello

Differential Revision: D96958741

fbshipit-source-id: 065f10c319dfe08c0a0c2c7cfe3b9a4407934afe

Author: DinoV

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

@@ -3,8 +3,10 @@
 
 import ast
 import sys
+from ast import AST
 from typing import Callable, Iterable
 
+from cinderx.compiler.errors import TypedSyntaxError
 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
@@ -56,6 +58,39 @@ def get_flags(
             override_flags
         )
 
+    def add_module(
+        self,
+        name: str,
+        filename: str,
+        tree: AST,
+        source: str | bytes | ast.Module | ast.Expression | ast.Interactive,
+        optimize: int,
+    ) -> ast.Module:
+        from cinderx.compiler.optimizer import AstOptimizer
+        from cinderx.compiler.static.declaration_visitor import DeclarationVisitor
+
+        optimized = AstOptimizer(optimize=optimize > 0).visit(tree)
+        assert isinstance(optimized, ast.Module)
+        tree = optimized
+
+        self.ast_cache[source] = tree
+
+        validate_classes = not self.decl_visitors
+        decl_visit = DeclarationVisitor(name, filename, self, optimize)
+        self.decl_visitors.append(decl_visit)
+        decl_visit.visit(tree)
+        decl_visit.finish_bind()
+
+        if validate_classes:
+            while self.decl_visitors:
+                decl_visit = self.decl_visitors.popleft()
+                try:
+                    decl_visit.module.validate_overrides()
+                except TypedSyntaxError:
+                    pass
+
+        return tree
+
     # pyre-ignore[14]: Pyre thinks the `compiler: Compiler` argument is inconsistent
     def make_type_binder(
         self,

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

@@ -254,7 +254,6 @@ def resolve_classname(
                 return result.inexact_type()
             elif isinstance(result, Value):
                 return result.klass
-            return None
 
         # Try well-known types (e.g. __static__.int64)
         well_known = type_env.name_to_type.get(qname)
@@ -266,7 +265,7 @@ def resolve_classname(
             builtins = modules.get("builtins")
             if builtins is not None:
                 # pyre-fixme[61]: `parts` is undefined, or not always defined.
-                result = builtins.get_child(parts[0], "builtins")
+                result = builtins.get_child(qname, "builtins")
                 if isinstance(result, Class):
                     return result
 

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

@@ -5,16 +5,29 @@
 from __future__ import annotations
 
 import ast
-from ast import AST, Name, Return
+from ast import AST, Attribute, Call, Compare, Constant, Expr, Name, Return
+from collections.abc import Sequence
 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 TerminalKind, TypeBinder
-from .types import CInstance, Class, CType
+from .type_binder import (
+    PRESERVE_REFINED_FIELDS,
+    PreserveRefinedFields,
+    TerminalKind,
+    TypeBinder,
+)
+from .types import (
+    # CInstance,
+    # Class,
+    # CType,
+    Dataclass,
+    DataclassField,
+    ModuleInstance,
+    TypeDescr,
+)
 
 if TYPE_CHECKING:
     from .compiler import Compiler
@@ -44,29 +57,105 @@ def __init__(
         self._type_info = type_info
 
     def visit(self, node: AST, *args: object) -> NarrowingEffect | None:
-        ret = super().visit(node, *args)
         if isinstance(node, ast.expr) and self._type_info is not None:
-            # If the parent visitor already resolved to a C primitive type
-            # (e.g. `x: int64 = 0` promotes 0 to int64, or a Name referencing
-            # a CType like int64), don't override with the pyrefly-inferred type.
-            existing_type = self.get_type(node)
-            if isinstance(existing_type, (CInstance, CType)):
-                return ret
-
+            ret = super().generic_visit(node, *args)
             # pyre-fixme[16]: Optional type has no attribute `lookup`.
             declared_type = self._type_info.lookup(node, self.modules, self.type_env)
 
             if declared_type is None:
                 declared_type = self.type_env.dynamic.instance
-
+                if isinstance(node, (ast.List, ast.ListComp)):
+                    declared_type = self.type_env.list.instance
+                elif isinstance(node, (ast.Dict, ast.DictComp)):
+                    declared_type = self.type_env.dict.instance
             self.set_type(node, declared_type)
-            if isinstance(node, Name) and isinstance(node.ctx, ast.Store):
-                try:
-                    self.declare_local(node.id, declared_type)
-                except TypedSyntaxError:
-                    pass  # already declared, just update the type
+
+            if isinstance(node, Compare):
+                for op in node.ops:
+                    self.set_type(op, self.type_env.DYNAMIC)
+
+            # Name: set PreserveRefinedFields (always), declare locals for
+            # Store context, and set TypeDescr for module-level names so
+            # that bind_call can emit direct invocations.  When pyrefly
+            # doesn't resolve the type, fall back to the module table so
+            # CinderX-specific types (e.g. ModuleInstance) are preserved.
+            elif isinstance(node, Name):
+                self.set_node_data(node, PreserveRefinedFields, PRESERVE_REFINED_FIELDS)
+                if isinstance(node.ctx, ast.Store):
+                    try:
+                        self.declare_local(node.id, declared_type)
+                    except TypedSyntaxError:
+                        pass  # already declared, just update the type
+                mod_typ, descr = self.module.resolve_name_with_descr(
+                    node.id, self.context_qualname
+                )
+                if descr is not None:
+                    self.set_node_data(node, TypeDescr, descr)
+                if (
+                    mod_typ is not None
+                    and declared_type is self.type_env.dynamic.instance
+                ):
+                    self.set_type(node, mod_typ)
+
+            # Attribute: when the base is a ModuleInstance, set TypeDescr
+            # for direct access and call bind_attr to resolve from the
+            # CinderX module table — this ensures CinderX-specific types
+            # (e.g. DataclassFieldFunction, DataclassDecorator) are used.
+            # Set PreserveRefinedFields when the attribute is refinable.
+            elif isinstance(node, Attribute):
+                base = self.get_type(node.value)
+                if isinstance(base, ModuleInstance):
+                    self.set_node_data(
+                        node, TypeDescr, ((base.module_name,), node.attr)
+                    )
+                    # Always call bind_attr for module attributes so that
+                    # CinderX-specific types (e.g. DataclassFieldFunction,
+                    # DataclassDecorator) are used regardless of what
+                    # pyrefly resolved.
+                    base.bind_attr(node, self, None)
+                if self.is_refinable(node):
+                    self.set_node_data(
+                        node, PreserveRefinedFields, PRESERVE_REFINED_FIELDS
+                    )
+
+            # Call: invoke bind_call on the func's type to populate
+            # ArgMapping (and ClassCallInfo for class instantiation).
+            elif isinstance(node, Call):
+                self.get_type(node.func).bind_call(node, self, None)
+                # When pyrefly compiles modules that aren't normally
+                # statically compiled, dataclasses.field() may resolve
+                # to DataclassField via the CinderX module table.
+                # In non-Dataclass classes (e.g. decorated with
+                # @deprecated on top of @dataclass), this causes type
+                # errors in visitAnnAssign.  Reset to dynamic and mark
+                # the class for non-static compilation in that case.
+                if isinstance(self.get_type(node), DataclassField) and not (
+                    isinstance(self.scope, ast.ClassDef)
+                    and isinstance(self.get_type(self.scope), Dataclass)
+                ):
+                    self.set_type(node, declared_type)
+                    if isinstance(self.scope, ast.ClassDef):
+                        self.module.compile_non_static.add(self.scope)
+        else:
+            return super().visit(node, *args)
+
         return ret
 
+    def visit_check_terminal(self, nodes: Sequence[ast.stmt]) -> TerminalKind:
+        # Treat a body consisting of just `...` (Ellipsis) as a terminal
+        # statement, so that stub functions like Protocol methods don't
+        # trigger "can implicitly return None" errors.
+        if (
+            len(nodes) == 1
+            and isinstance(nodes[0], Expr)
+            and isinstance(nodes[0].value, Constant)
+            and nodes[0].value.value is ...
+        ):
+            self.visit(nodes[0])
+            self.set_terminal_kind(nodes[0], TerminalKind.RaiseOrReturn)
+            return TerminalKind.RaiseOrReturn
+        return super().visit_check_terminal(nodes)
+
     def visitReturn(self, node: Return) -> None:
         self.set_terminal_kind(node, TerminalKind.RaiseOrReturn)
         if node.value is not None:

cinderx/PythonLib/test_cinderx/test_compiler/test_static/pyreflytests/method_call.test.py

@@ -6,6 +6,5 @@
 def verify(test: PyreBinderTests, code: CodeType) -> None:
     """Verify that the compiled code includes INVOKE_METHOD."""
 
-    # TODO: Enable this when things are working
     g = test.find_code(code, "g")
-    test.assertInBytecode(g, "INVOKE_METHOD")
+    test.assertInBytecode(g, "INVOKE_FUNCTION")

cinderx/PythonLib/test_cinderx/test_compiler/test_static/pyreflytests/property.test.py

@@ -7,4 +7,4 @@ def verify(test: PyreBinderTests, code: CodeType) -> None:
     """Verify that the compiled code includes INVOKE_METHOD."""
 
     g = test.find_code(code, "g")
-    test.assertInBytecode(g, "INVOKE_METHOD")
+    test.assertInBytecode(g, "INVOKE_FUNCTION")