From 539513a5ff41954b0d0d5b0cd25b1cb03a10122b Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Thu, 17 Oct 2024 02:18:47 -0700
Subject: [PATCH] Fix iteration over union (when self type is involved)
 (#17976)

Fixes https://github.com/python/mypy/issues/17945
---
 mypy/checkexpr.py                  |  1 +
 test-data/unit/check-selftype.test | 16 ++++++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 8d1833a772e2f..b06aaa8f89f54 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -3814,6 +3814,7 @@ def check_method_call_by_name(
             is_operator=True,
             msg=self.msg,
             original_type=original_type,
+            self_type=base_type,
             chk=self.chk,
             in_literal_context=self.is_literal_context(),
         )
diff --git a/test-data/unit/check-selftype.test b/test-data/unit/check-selftype.test
index 9601852ef823b..1a088fe05092a 100644
--- a/test-data/unit/check-selftype.test
+++ b/test-data/unit/check-selftype.test
@@ -2160,3 +2160,19 @@ class MyProtocol(Protocol):
 
 def test() -> None: ...
 value: MyProtocol = test
+
+[case testSelfTypeUnionIter]
+from typing import Self, Iterator, Generic, TypeVar, Union
+
+T = TypeVar("T")
+
+class range(Generic[T]):
+    def __iter__(self) -> Self: ...
+    def __next__(self) -> T: ...
+
+class count:
+    def __iter__(self) -> Iterator[int]: ...
+
+def foo(x: Union[range[int], count]) -> None:
+    for item in x:
+        reveal_type(item)  # N: Revealed type is "builtins.int"