Mypy now treats classes with __getitem__ as iterable

Author: sobolevn

mypy/checker.py

@@ -3429,9 +3429,9 @@ def type_is_iterable(self, type: Type) -> bool:
         type = get_proper_type(type)
         if isinstance(type, CallableType) and type.is_type_obj():
             type = type.fallback
-        return is_subtype(
-            type, self.named_generic_type("typing.Iterable", [AnyType(TypeOfAny.special_form)])
-        )
+        with self.msg.filter_errors() as iter_errors:
+            self.analyze_iterable_item_type(TempNode(type))
+        return not iter_errors.has_new_errors()
 
     def check_multi_assignment_from_iterable(
         self,
@@ -4247,15 +4247,36 @@ def analyze_iterable_item_type(self, expr: Expression) -> tuple[Type, Type]:
         """Analyse iterable expression and return iterator and iterator item types."""
         echk = self.expr_checker
         iterable = get_proper_type(echk.accept(expr))
-        iterator = echk.check_method_call_by_name("__iter__", iterable, [], [], expr)[0]
 
+        # We first try to find `__iter__` magic method.
+        # If it is present, we go on with it.
+        # But, python also support iterables with just `__getitem__(index) -> Any` defined.
+        # So, we check it in case `__iter__` is missing.
+        with self.msg.filter_errors(save_filtered_errors=True) as iter_errors:
+            # We save original error to show it later if `__getitem__` is also missing.
+            iterator = echk.check_method_call_by_name("__iter__", iterable, [], [], expr)[0]
+        if iter_errors.has_new_errors():
+            # `__iter__` is missing, try `__getattr__`:
+            arg = self.temp_node(AnyType(TypeOfAny.implementation_artifact), expr)
+            with self.msg.filter_errors() as getitem_errors:
+                getitem_type = echk.check_method_call_by_name(
+                    "__getitem__", iterable, [arg], [nodes.ARG_POS], expr
+                )[0]
+            if getitem_errors.has_new_errors():  # Both are missing.
+                self.msg.add_errors(iter_errors.filtered_errors())
+                return AnyType(TypeOfAny.from_error), AnyType(TypeOfAny.from_error)
+            else:
+                # We found just `__getitem__`, it does not follow the same
+                # semantics as `__iter__`, so: just return what we found.
+                return self.named_generic_type("typing.Iterator", [getitem_type]), getitem_type
+
+        # We found `__iter__`, let's analyze its return type:
         if isinstance(iterable, TupleType):
             joined: Type = UninhabitedType()
             for item in iterable.items:
                 joined = join_types(joined, item)
             return iterator, joined
-        else:
-            # Non-tuple iterable.
+        else:  # Non-tuple iterable.
             return iterator, echk.check_method_call_by_name("__next__", iterator, [], [], expr)[0]
 
     def analyze_container_item_type(self, typ: Type) -> Type | None:
@@ -5986,15 +6007,9 @@ def iterable_item_type(self, instance: Instance) -> Type:
             # This relies on 'map_instance_to_supertype' returning 'Iterable[Any]'
             # in case there is no explicit base class.
             return item_type
-        # Try also structural typing.
-        iter_type = get_proper_type(find_member("__iter__", instance, instance, is_operator=True))
-        if iter_type and isinstance(iter_type, CallableType):
-            ret_type = get_proper_type(iter_type.ret_type)
-            if isinstance(ret_type, Instance):
-                iterator = map_instance_to_supertype(
-                    ret_type, self.lookup_typeinfo("typing.Iterator")
-                )
-                item_type = iterator.args[0]
+
+        # Try also structural typing: including `__iter__` and `__getitem__`.
+        _, item_type = self.analyze_iterable_item_type(TempNode(instance))
         return item_type
 
     def function_type(self, func: FuncBase) -> FunctionLike:

mypy/checkexpr.py

@@ -2998,13 +2998,14 @@ def check_method_call_by_name(
             method,
             base_type,
             context,
-            False,
-            False,
-            True,
-            self.msg,
+            is_lvalue=False,
+            is_super=False,
+            is_operator=True,
+            msg=self.msg,
             original_type=original_type,
             chk=self.chk,
             in_literal_context=self.is_literal_context(),
+            suggest_awaitable=False,
         )
         return self.check_method_call(method, base_type, method_type, args, arg_kinds, context)
 
@@ -4806,13 +4807,8 @@ def visit_yield_from_expr(self, e: YieldFromExpr, allow_none_return: bool = Fals
             if is_async_def(subexpr_type) and not has_coroutine_decorator(return_type):
                 self.chk.msg.yield_from_invalid_operand_type(subexpr_type, e)
 
-            any_type = AnyType(TypeOfAny.special_form)
-            generic_generator_type = self.chk.named_generic_type(
-                "typing.Generator", [any_type, any_type, any_type]
-            )
-            iter_type, _ = self.check_method_call_by_name(
-                "__iter__", subexpr_type, [], [], context=generic_generator_type
-            )
+            iter_type, _ = self.chk.analyze_iterable_item_type(TempNode(subexpr_type))
+            iter_type = get_proper_type(iter_type)
         else:
             if not (is_async_def(subexpr_type) and has_coroutine_decorator(return_type)):
                 self.chk.msg.yield_from_invalid_operand_type(subexpr_type, e)

mypy/checkmember.py

@@ -90,6 +90,7 @@ def __init__(
         chk: mypy.checker.TypeChecker,
         self_type: Type | None,
         module_symbol_table: SymbolTable | None = None,
+        suggest_awaitable: bool = True,
     ) -> None:
         self.is_lvalue = is_lvalue
         self.is_super = is_super
@@ -100,6 +101,7 @@ def __init__(
         self.msg = msg
         self.chk = chk
         self.module_symbol_table = module_symbol_table
+        self.suggest_awaitable = suggest_awaitable
 
     def named_type(self, name: str) -> Instance:
         return self.chk.named_type(name)
@@ -149,6 +151,7 @@ def analyze_member_access(
     in_literal_context: bool = False,
     self_type: Type | None = None,
     module_symbol_table: SymbolTable | None = None,
+    suggest_awaitable: bool = True,
 ) -> Type:
     """Return the type of attribute 'name' of 'typ'.
 
@@ -183,6 +186,7 @@ def analyze_member_access(
         chk=chk,
         self_type=self_type,
         module_symbol_table=module_symbol_table,
+        suggest_awaitable=suggest_awaitable,
     )
     result = _analyze_member_access(name, typ, mx, override_info)
     possible_literal = get_proper_type(result)
@@ -260,7 +264,7 @@ def report_missing_attribute(
     override_info: TypeInfo | None = None,
 ) -> Type:
     res_type = mx.msg.has_no_attr(original_type, typ, name, mx.context, mx.module_symbol_table)
-    if may_be_awaitable_attribute(name, typ, mx, override_info):
+    if mx.suggest_awaitable and may_be_awaitable_attribute(name, typ, mx, override_info):
         mx.msg.possible_missing_await(mx.context)
     return res_type
 

mypy/constraints.py

@@ -750,8 +750,7 @@ def infer_constraints_from_protocol_members(
         """
         res = []
         for member in protocol.type.protocol_members:
-            inst = mypy.subtypes.find_member(member, instance, subtype)
-            temp = mypy.subtypes.find_member(member, template, subtype)
+            inst, temp = mypy.subtypes.find_members(member, instance, template, subtype)
             if inst is None or temp is None:
                 return []  # See #11020
             # The above is safe since at this point we know that 'instance' is a subtype

mypy/subtypes.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
-from typing import Any, Callable, Iterator, List, TypeVar, cast
+from typing import Any, Callable, Iterator, List, Tuple, TypeVar, cast
 from typing_extensions import Final, TypeAlias as _TypeAlias
 
 import mypy.applytype
@@ -14,6 +14,7 @@
 # Circular import; done in the function instead.
 # import mypy.solve
 from mypy.nodes import (
+    ARG_POS,
     ARG_STAR,
     ARG_STAR2,
     CONTRAVARIANT,
@@ -956,9 +957,10 @@ def f(self) -> A: ...
             ignore_names = member != "__call__"  # __call__ can be passed kwargs
             # The third argument below indicates to what self type is bound.
             # We always bind self to the subtype. (Similarly to nominal types).
-            supertype = get_proper_type(find_member(member, right, left))
+            supertype, subtype = find_members(member, right, left, left)
+            supertype = get_proper_type(supertype)
             assert supertype is not None
-            subtype = get_proper_type(find_member(member, left, left))
+            subtype = get_proper_type(subtype)
             # Useful for debugging:
             # print(member, 'of', left, 'has type', subtype)
             # print(member, 'of', right, 'has type', supertype)
@@ -1012,6 +1014,56 @@ def f(self) -> A: ...
     return True
 
 
+def find_members(
+    name: str, supertype: Instance, subtype: Instance, context: Type
+) -> Tuple[Type | None, Type | None]:
+    """Find types of member by name for two instances.
+
+    We do it with respect to some special cases, like `Iterable` and `__geitem__`.
+    """
+    if name == "__iter__":
+        # So, this is a special case: old-style iterbale protocol
+        # must be supported even without explicit `__iter__` method.
+        # Because all types with `__geitem__` defined have default `__iter__`
+        # implementation. See #2220
+        # First, we need to find which is one actually `Iterable`:
+        if is_named_instance(supertype, "typing.Iterable"):
+            left, right = _iterable_special_member(supertype, subtype, context)
+            if left is not None and right is not None:
+                return left, right
+        elif is_named_instance(subtype, "typing.Iterable"):
+            left, right = _iterable_special_member(subtype, supertype, context)
+            if left is not None and right is not None:
+                return right, left
+
+    # This is not a special case.
+    # Falling back to regular `find_member` call:
+    return (find_member(name, supertype, context), find_member(name, subtype, context))
+
+
+def _iterable_special_member(
+    iterable: Instance, candidate: Instance, context: Type
+) -> Tuple[Type | None, Type | None]:
+    name = "__iter__"
+    iterable_method = get_proper_type(find_member(name, iterable, context))
+    candidate_method = get_proper_type(find_member("__getitem__", candidate, context))
+    if isinstance(iterable_method, CallableType) and isinstance(
+        (ret := get_proper_type(iterable_method.ret_type)), Instance
+    ):
+        # We need to transform
+        # `__iter__() -> Iterable[ret]` into
+        # `__getitem__(Any) -> ret`
+        iterable_method = iterable_method.copy_modified(
+            arg_names=[None],
+            arg_types=[AnyType(TypeOfAny.implementation_artifact)],
+            arg_kinds=[ARG_POS],
+            ret_type=ret.args[0],
+            name="__getitem__",
+        )
+        return (iterable_method, candidate_method)
+    return None, None
+
+
 def find_member(
     name: str, itype: Instance, subtype: Type, is_operator: bool = False
 ) -> Type | None: