From 34d86038bab0b79251019fd1cbb438b7aef1a592 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 770ba4bd352f..ffd72900bb2e 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -3805,6 +3805,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 9601852ef823..1a088fe05092 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"