Merge branch 'master' into patch-3

Author: BobTheBuidler

.github/workflows/mypy_primer.yml

@@ -28,7 +28,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        shard-index: [0, 1, 2, 3, 4]
+        shard-index: [0, 1, 2, 3, 4, 5]
       fail-fast: false
     timeout-minutes: 60
     steps:
@@ -63,7 +63,7 @@ jobs:
             mypy_primer \
             --repo mypy_to_test \
             --new $GITHUB_SHA --old base_commit \
-            --num-shards 5 --shard-index ${{ matrix.shard-index }} \
+            --num-shards 6 --shard-index ${{ matrix.shard-index }} \
             --debug \
             --additional-flags="--debug-serialize" \
             --output concise \

mypy/checker.py

@@ -3138,7 +3138,7 @@ def check_assignment(
                         else:
                             self.check_getattr_method(signature, lvalue, name)
 
-                if name == "__slots__":
+                if name == "__slots__" and self.scope.active_class() is not None:
                     typ = lvalue_type or self.expr_checker.accept(rvalue)
                     self.check_slots_definition(typ, lvalue)
                 if name == "__match_args__" and inferred is not None:
@@ -3317,6 +3317,12 @@ def get_variable_type_context(self, inferred: Var, rvalue: Expression) -> Type |
                     type_contexts.append(base_type)
         # Use most derived supertype as type context if available.
         if not type_contexts:
+            if inferred.name == "__slots__" and self.scope.active_class() is not None:
+                str_type = self.named_type("builtins.str")
+                return self.named_generic_type("typing.Iterable", [str_type])
+            if inferred.name == "__all__" and self.scope.is_top_level():
+                str_type = self.named_type("builtins.str")
+                return self.named_generic_type("typing.Sequence", [str_type])
             return None
         candidate = type_contexts[0]
         for other in type_contexts:

mypy/checker_shared.py

@@ -334,6 +334,10 @@ def current_self_type(self) -> Instance | TupleType | None:
                 return fill_typevars(item)
         return None
 
+    def is_top_level(self) -> bool:
+        """Is current scope top-level (no classes or functions)?"""
+        return len(self.stack) == 1
+
     @contextmanager
     def push_function(self, item: FuncItem) -> Iterator[None]:
         self.stack.append(item)

mypy/plugins/enums.py

@@ -17,10 +17,12 @@
 from typing import TypeVar, cast
 
 import mypy.plugin  # To avoid circular imports.
-from mypy.nodes import TypeInfo
+from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo, Var
 from mypy.subtypes import is_equivalent
 from mypy.typeops import fixup_partial_type, make_simplified_union
 from mypy.types import (
+    ELLIPSIS_TYPE_NAMES,
     CallableType,
     Instance,
     LiteralType,
@@ -79,6 +81,19 @@ def _infer_value_type_with_auto_fallback(
     if proper_type is None:
         return None
     proper_type = get_proper_type(fixup_partial_type(proper_type))
+    # Enums in stubs may have ... instead of actual values. If `_value_` is annotated
+    # (manually or inherited from IntEnum, for example), it is a more reasonable guess
+    # than literal ellipsis type.
+    if (
+        _is_defined_in_stub(ctx)
+        and isinstance(proper_type, Instance)
+        and proper_type.type.fullname in ELLIPSIS_TYPE_NAMES
+        and isinstance(ctx.type, Instance)
+    ):
+        value_type = ctx.type.type.get("_value_")
+        if value_type is not None and isinstance(var := value_type.node, Var):
+            return var.type
+        return proper_type
     if not (isinstance(proper_type, Instance) and proper_type.type.fullname == "enum.auto"):
         if is_named_instance(proper_type, "enum.member") and proper_type.args:
             return proper_type.args[0]
@@ -106,6 +121,11 @@ def _infer_value_type_with_auto_fallback(
     return ctx.default_attr_type
 
 
+def _is_defined_in_stub(ctx: mypy.plugin.AttributeContext) -> bool:
+    assert isinstance(ctx.api, TypeChecker)
+    return isinstance(ctx.type, Instance) and ctx.api.modules[ctx.type.type.module_name].is_stub
+
+
 def _implements_new(info: TypeInfo) -> bool:
     """Check whether __new__ comes from enum.Enum or was implemented in a
     subclass. In the latter case, we must infer Any as long as mypy can't infer

mypy/stubtest.py

@@ -1149,7 +1149,7 @@ def verify_var(
             proper_type = mypy.types.get_proper_type(stub.type)
             if (
                 isinstance(proper_type, mypy.types.Instance)
-                and proper_type.type.fullname == "builtins.ellipsis"
+                and proper_type.type.fullname in mypy.types.ELLIPSIS_TYPE_NAMES
             ):
                 should_error = False
 

mypy/types.py

@@ -181,6 +181,8 @@
 # Supported @override decorator names.
 OVERRIDE_DECORATOR_NAMES: Final = ("typing.override", "typing_extensions.override")
 
+ELLIPSIS_TYPE_NAMES: Final = ("builtins.ellipsis", "types.EllipsisType")
+
 # A placeholder used for Bogus[...] parameters
 _dummy: Final[Any] = object()
 
@@ -1574,7 +1576,7 @@ def is_singleton_type(self) -> bool:
         return (
             self.type.is_enum
             and len(self.type.enum_members) == 1
-            or self.type.fullname in {"builtins.ellipsis", "types.EllipsisType"}
+            or self.type.fullname in ELLIPSIS_TYPE_NAMES
         )
 
 

mypyc/irbuild/for_helpers.py

@@ -24,19 +24,23 @@
     TypeAlias,
 )
 from mypyc.ir.ops import (
+    ERR_NEVER,
     BasicBlock,
     Branch,
     Integer,
     IntOp,
     LoadAddress,
+    LoadErrorValue,
     LoadMem,
+    MethodCall,
     RaiseStandardError,
     Register,
     TupleGet,
     TupleSet,
     Value,
 )
 from mypyc.ir.rtypes import (
+    RInstance,
     RTuple,
     RType,
     bool_rprimitive,
@@ -48,10 +52,13 @@
     is_short_int_rprimitive,
     is_str_rprimitive,
     is_tuple_rprimitive,
+    object_pointer_rprimitive,
+    object_rprimitive,
     pointer_rprimitive,
     short_int_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
 from mypyc.primitives.dict_ops import (
     dict_check_size_op,
@@ -62,7 +69,7 @@
     dict_next_value_op,
     dict_value_iter_op,
 )
-from mypyc.primitives.exc_ops import no_err_occurred_op
+from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op
 from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op
 from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
 from mypyc.primitives.misc_ops import stop_async_iteration_op
@@ -511,7 +518,15 @@ def make_for_loop_generator(
     # Default to a generic for loop.
     if iterable_expr_reg is None:
         iterable_expr_reg = builder.accept(expr)
-    for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
+
+    it = iterable_expr_reg.type
+    for_obj: ForNativeGenerator | ForIterable
+    if isinstance(it, RInstance) and it.class_ir.has_method(GENERATOR_HELPER_NAME):
+        # Directly call generator object methods if iterating over a native generator.
+        for_obj = ForNativeGenerator(builder, index, body_block, loop_exit, line, nested)
+    else:
+        # Generic implementation that works of arbitrary iterables.
+        for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested)
     item_type = builder._analyze_iterable_item_type(expr)
     item_rtype = builder.type_to_rtype(item_type)
     for_obj.init(iterable_expr_reg, item_rtype)
@@ -623,6 +638,63 @@ def gen_cleanup(self) -> None:
         self.builder.call_c(no_err_occurred_op, [], self.line)
 
 
+class ForNativeGenerator(ForGenerator):
+    """Generate IR for a for loop over a native generator."""
+
+    def need_cleanup(self) -> bool:
+        # Create a new cleanup block for when the loop is finished.
+        return True
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        # Define target to contains the generator expression. It's also the iterator.
+        # If we are inside a generator function, spill these into the environment class.
+        builder = self.builder
+        self.iter_target = builder.maybe_spill(expr_reg)
+        self.target_type = target_type
+
+    def gen_condition(self) -> None:
+        builder = self.builder
+        line = self.line
+        self.return_value = Register(object_rprimitive)
+        err = builder.add(LoadErrorValue(object_rprimitive, undefines=True))
+        builder.assign(self.return_value, err, line)
+
+        # Call generated generator helper method, passing a PyObject ** as the final
+        # argument that will be used to store the return value in the return value
+        # register. We ignore the return value but the presence of a return value
+        # indicates that the generator has finished. This is faster than raising
+        # and catching StopIteration, which is the non-native way of doing this.
+        ptr = builder.add(LoadAddress(object_pointer_rprimitive, self.return_value))
+        nn = builder.none_object()
+        helper_call = MethodCall(
+            builder.read(self.iter_target), GENERATOR_HELPER_NAME, [nn, nn, nn, nn, ptr], line
+        )
+        # We provide custom handling for error values.
+        helper_call.error_kind = ERR_NEVER
+
+        self.next_reg = builder.add(helper_call)
+        builder.add(Branch(self.next_reg, self.loop_exit, self.body_block, Branch.IS_ERROR))
+
+    def begin_body(self) -> None:
+        # Assign the value obtained from the generator helper method to the
+        # lvalue so that it can be referenced by code in the body of the loop.
+        builder = self.builder
+        line = self.line
+        # We unbox here so that iterating with tuple unpacking generates a tuple based
+        # unpack instead of an iterator based one.
+        next_reg = builder.coerce(self.next_reg, self.target_type, line)
+        builder.assign(builder.get_assignment_target(self.index), next_reg, line)
+
+    def gen_step(self) -> None:
+        # Nothing to do here, since we get the next item as part of gen_condition().
+        pass
+
+    def gen_cleanup(self) -> None:
+        # If return value is NULL (it wasn't assigned to by the generator helper method),
+        # an exception was raised that we need to propagate.
+        self.builder.primitive_op(propagate_if_error_op, [self.return_value], self.line)
+
+
 class ForAsyncIterable(ForGenerator):
     """Generate IR for an async for loop."""
 

mypyc/irbuild/prepare.py

@@ -15,7 +15,7 @@
 
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import NamedTuple
+from typing import Final, NamedTuple
 
 from mypy.build import Graph
 from mypy.nodes import (
@@ -71,7 +71,7 @@
 from mypyc.options import CompilerOptions
 from mypyc.sametype import is_same_type
 
-GENERATOR_HELPER_NAME = "__mypyc_generator_helper__"
+GENERATOR_HELPER_NAME: Final = "__mypyc_generator_helper__"
 
 
 def build_type_map(