1 files changed, 46 insertions, 0 deletions
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index 9606d5b..ec2b68a 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -10,6 +10,7 @@ import collections.abc
 import functools
 import contextlib
 import builtins
+import traceback
 
 # Test result of triple loop (too big to inline)
 TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
@@ -1143,6 +1144,51 @@ class TestCase(unittest.TestCase):
             self.assertRaises(TypeError, iter, typ())
         self.assertRaises(ZeroDivisionError, iter, BadIterableClass())
 
+    def test_exception_locations(self):
+        # The location of an exception raised from __init__ or
+        # __next__ should should be the iterator expression
+
+        class Iter:
+            def __init__(self, init_raises=False, next_raises=False):
+                if init_raises:
+                    1/0
+                self.next_raises = next_raises
+
+            def __next__(self):
+                if self.next_raises:
+                    1/0
+
+            def __iter__(self):
+                return self
+
+        def init_raises():
+            try:
+                for x in Iter(init_raises=True):
+                    pass
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                for x in Iter(next_raises=True):
+                    pass
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "Iter(init_raises=True)"),
+                               (next_raises, "Iter(next_raises=True)"),
+                              ]:
+            with self.subTest(func):
+                exc = func()
+                f = traceback.extract_tb(exc.__traceback__)[0]
+                indent = 16
+                co = func.__code__
+                self.assertEqual(f.lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
+                                 expected)
+
+
 
 if __name__ == "__main__":
     unittest.main()