Fix type narrowing

Author: michaelm-openai

mypy/checker.py

@@ -6436,6 +6436,10 @@ def find_isinstance_check_helper(
                 # If the callee is a RefExpr, extract TypeGuard/TypeIs directly.
                 if isinstance(node.callee, RefExpr):
                     type_is, type_guard = node.callee.type_is, node.callee.type_guard
+                if type_guard is not None:
+                    type_guard = self.expand_narrowed_type(type_guard)
+                if type_is is not None:
+                    type_is = self.expand_narrowed_type(type_is)
                 if type_guard is not None or type_is is not None:
                     # TODO: Follow *args, **kwargs
                     if node.arg_kinds[0] != nodes.ARG_POS:
@@ -7926,7 +7930,7 @@ def conditional_types_with_intersection(
                 for types, reason in errors:
                     self.msg.impossible_intersection(types, reason, ctx)
             return UninhabitedType(), expr_type
-        new_yes_type = make_simplified_union(out)
+        new_yes_type: Type = make_simplified_union(out)
         return new_yes_type, expr_type
 
     def is_writable_attribute(self, node: Node) -> bool:
@@ -7985,8 +7989,31 @@ def get_isinstance_type(self, expr: Expression) -> list[TypeRange] | None:
                 types.append(TypeRange(typ, is_upper_bound=False))
             else:  # we didn't see an actual type, but rather a variable with unknown value
                 return None
+        return self.expand_isinstance_type_ranges(types)
+
+    def expand_isinstance_type_ranges(self, types: list[TypeRange]) -> list[TypeRange]:
+        if disallow_str_iteration_state.disallow_str_iteration:
+            str_type = self.named_type("builtins.str")
+            return types + [
+                TypeRange(str_type, is_upper_bound=type_range.is_upper_bound)
+                for type_range in types
+                if self._is_str_iteration_protocol_for_narrowing(type_range.item)
+            ]
         return types
 
+    def _is_str_iteration_protocol_for_narrowing(self, typ: Type) -> bool:
+        proper = get_proper_type(typ)
+        return isinstance(proper, Instance) and proper.type.fullname in STR_ITERATION_PROTOCOL_BASES
+
+    def expand_narrowed_type(self, typ: Type) -> Type:
+        if disallow_str_iteration_state.disallow_str_iteration:
+            proper = get_proper_type(typ)
+            if isinstance(proper, UnionType):
+                return make_simplified_union([self.expand_narrowed_type(item) for item in proper.items])
+            if self._is_str_iteration_protocol_for_narrowing(proper):
+                return make_simplified_union([typ, self.named_type("builtins.str")])
+        return typ
+
     def is_literal_enum(self, n: Expression) -> bool:
         """Returns true if this expression (with the given type context) is an Enum literal.
 

test-data/unit/check-flags.test

@@ -2452,9 +2452,9 @@ f(memoryview(b"asdf"))  # E: Argument 1 to "f" has incompatible type "memoryview
 [builtins fixtures/primitives.pyi]
 
 [case testDisallowStrIteration]
-# flags: --disallow-str-iteration
+# flags: --disallow-str-iteration --warn-unreachable
 from abc import abstractmethod
-from typing import Collection, Container, Iterable, Mapping, Protocol, Sequence, SupportsIndex, TypeVar, Union
+from typing import Collection, Container, Hashable, Iterable, Mapping, Protocol, Sequence, SupportsIndex, TypeGuard, TypeIs, TypeVar, Union, runtime_checkable
 
 def takes_str(x: str):
     for ch in x:  # E: Iterating over "str" is disallowed  # N: This is because --disallow-str-iteration is enabled
@@ -2476,7 +2476,7 @@ seq: Sequence[str] = s  # E: Incompatible types in assignment (expression has ty
 iterable: Iterable[str] = s  # E: Incompatible types in assignment (expression has type "str", variable has type "Iterable[str]")
 collection: Collection[str] = s  # E: Incompatible types in assignment (expression has type "str", variable has type "Collection[str]")
 
-def takes_maybe_seq(x: "str | Sequence[int]") -> None:
+def takes_maybe_seq(x: Union[str, Sequence[int]]) -> None:
     for ch in x:  # E: Iterating over "str | Sequence[int]" is disallowed  # N: This is because --disallow-str-iteration is enabled
         reveal_type(ch)  # N: Revealed type is "builtins.str | builtins.int"
 
@@ -2495,6 +2495,7 @@ def takes_str_subclass(x: StrSubclass):
     for ch in x:  # E: Iterating over "StrSubclass" is disallowed  # N: This is because --disallow-str-iteration is enabled
         reveal_type(ch)  # N: Revealed type is "builtins.str"
 
+@runtime_checkable
 class CollectionSubclass(Collection[_T_co], Protocol[_T_co]):
     @abstractmethod
     def __missing_impl__(self): ...
@@ -2512,12 +2513,52 @@ takes_collection_subclass(StrSubclass())  # E: Argument 1 to "takes_collection_s
 def dict_unpacking_unaffected_by_union_simplification(x: Mapping[str, Union[str, Sequence[str]]]) -> None:
     x = {**x}
 
-def narrowing(x: "str | Sequence[str]"):
+def narrowing_str(x: Union[str, Sequence[str]]):
     if isinstance(x, str):
         reveal_type(x)  # N: Revealed type is "builtins.str"
     else:
         reveal_type(x)  # N: Revealed type is "typing.Sequence[builtins.str]"
 
+def narrowing_iterable(x: Union[str, Iterable[str], Iterable[str]]):
+    if isinstance(x, Iterable):
+        reveal_type(x)  # N: Revealed type is "builtins.str | typing.Iterable[builtins.str]"
+    else:
+        reveal_type(x)  # E: Statement is unreachable
+
+def is_iterable_guard(x: object) -> TypeGuard[Iterable[str]]: ...
+
+def is_iterable_is(x: object) -> TypeIs[Iterable[str]]: ...
+
+def narrowing_typeguard_iterable(x: Union[str, Iterable[str]]):
+    if is_iterable_guard(x):
+        reveal_type(x)  # N: Revealed type is "builtins.str | typing.Iterable[builtins.str]"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str | typing.Iterable[builtins.str]"
+
+def narrowing_typeis_iterable(x: Union[str, Iterable[str]]):
+    if is_iterable_is(x):
+        reveal_type(x)  # N: Revealed type is "builtins.str | typing.Iterable[builtins.str]"
+    else:
+        reveal_type(x)  # E: Statement is unreachable
+
+def narrowing_incompatible_collection_subclass(x: Union[str, CollectionSubclass[str]]):
+    if isinstance(x, CollectionSubclass):
+        reveal_type(x)  # N: Revealed type is "__main__.CollectionSubclass[builtins.str]"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.str"
+
+def narrowing_object(x: object):
+    if isinstance(x, Iterable):
+        reveal_type(x)  # N: Revealed type is "typing.Iterable[Any] | builtins.str"
+    else:
+        reveal_type(x)  # N: Revealed type is "builtins.object"
+
+def narrowing_type(x: Union[type[str], type[int]]):
+    if issubclass(x, Iterable):
+        reveal_type(x)  # N: Revealed type is "type[builtins.str]"
+    else:
+        reveal_type(x)  # N: Revealed type is "type[builtins.int]"
+
 Item = TypeVar("Item")
 def takes_generic_list(x: list[Item]) -> None: ...
 takes_generic_list(reveal_type([s, seq]))  # N: Revealed type is "builtins.list[builtins.object]"

test-data/unit/fixtures/str-iter.pyi

@@ -1,14 +1,15 @@
 # Builtins stub used in disallow-str-iteration tests.
 
 
-from typing import Generic, Iterator, Mapping, Sequence, SupportsIndex,TypeVar, overload
+from typing import Generic, Iterator, Mapping, Sequence, SupportsIndex, Type, TypeVar, overload
 
 _T = TypeVar("_T")
 _KT = TypeVar("_KT")
 _VT = TypeVar("_VT")
 
 class object:
     def __init__(self) -> None: pass
+    def __hash__(self) -> int: pass
 
 class type: pass
 class int: pass
@@ -21,6 +22,7 @@ class str(Sequence[str]):
     def __len__(self) -> int: pass
     def __contains__(self, item: object) -> bool: pass
     def __getitem__(self, key: SupportsIndex | slice, /) -> str: pass
+    def __hash__(self) -> int: pass
 
 class list(Sequence[_T], Generic[_T]):
     def __iter__(self) -> Iterator[_T]: pass
@@ -47,3 +49,4 @@ class dict(Mapping[_KT, _VT], Generic[_KT, _VT]):
     def __getitem__(self, key: _KT) -> _VT: pass
 
 def isinstance(x: object, t: type) -> bool: pass
+def issubclass(x: type, t: type) -> bool: pass

test-data/unit/fixtures/typing-str-iter.pyi

@@ -8,6 +8,9 @@ TypeVar = 0
 Generic = 0
 Protocol = 0
 Union = 0
+Type = 0
+TypeGuard = 0
+TypeIs = 0
 overload = 0
 
 _T = TypeVar("_T")
@@ -17,25 +20,35 @@ _KT_co = TypeVar("_KT_co", covariant=True)  # Key type covariant containers.
 _VT_co = TypeVar("_VT_co", covariant=True)  # Value type covariant containers.
 _TC = TypeVar("_TC", bound=type[object])
 
+@runtime_checkable
+class Hashable(Protocol, metaclass=ABCMeta):
+    @abstractmethod
+    def __hash__(self) -> int: pass
+
+@runtime_checkable
 class Iterable(Protocol[_T_co]):
     @abstractmethod
     def __iter__(self) -> Iterator[_T_co]: ...
 
+@runtime_checkable
 class Iterator(Iterable[_T_co], Protocol[_T_co]):
     @abstractmethod
     def __next__(self) -> _T_co: ...
     def __iter__(self) -> Iterator[_T_co]: ...
 
+@runtime_checkable
 class Container(Protocol[_T_co]):
     # This is generic more on vibes than anything else
     @abstractmethod
     def __contains__(self, x: object, /) -> bool: ...
 
+@runtime_checkable
 class Collection(Iterable[_T_co], Container[_T_co], Protocol[_T_co]):
     # Implement Sized (but don't have it as a base class).
     @abstractmethod
     def __len__(self) -> int: ...
 
+@runtime_checkable
 class SupportsIndex(Protocol, metaclass=ABCMeta):
     @abstractmethod
     def __index__(self) -> int: ...