Close #15153: Added inspect.getgeneratorlocals to simplify whitebox testing of generator state updates
diff --git a/Lib/inspect.py b/Lib/inspect.py
index dd2de64..074e1b4 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -1259,6 +1259,8 @@
raise AttributeError(attr)
+# ------------------------------------------------ generator introspection
+
GEN_CREATED = 'GEN_CREATED'
GEN_RUNNING = 'GEN_RUNNING'
GEN_SUSPENDED = 'GEN_SUSPENDED'
@@ -1282,6 +1284,22 @@
return GEN_SUSPENDED
+def getgeneratorlocals(generator):
+ """
+ Get the mapping of generator local variables to their current values.
+
+ A dict is returned, with the keys the local variable names and values the
+ bound values."""
+
+ if not isgenerator(generator):
+ raise TypeError("'{!r}' is not a Python generator".format(generator))
+
+ frame = getattr(generator, "gi_frame", None)
+ if frame is not None:
+ return generator.gi_frame.f_locals
+ else:
+ return {}
+
###############################################################################
### Function Signature Object (PEP 362)
###############################################################################
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 8327721..7f70342 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -1271,6 +1271,52 @@
self.assertIn(name, repr(state))
self.assertIn(name, str(state))
+ def test_getgeneratorlocals(self):
+ def each(lst, a=None):
+ b=(1, 2, 3)
+ for v in lst:
+ if v == 3:
+ c = 12
+ yield v
+
+ numbers = each([1, 2, 3])
+ self.assertEqual(inspect.getgeneratorlocals(numbers),
+ {'a': None, 'lst': [1, 2, 3]})
+ next(numbers)
+ self.assertEqual(inspect.getgeneratorlocals(numbers),
+ {'a': None, 'lst': [1, 2, 3], 'v': 1,
+ 'b': (1, 2, 3)})
+ next(numbers)
+ self.assertEqual(inspect.getgeneratorlocals(numbers),
+ {'a': None, 'lst': [1, 2, 3], 'v': 2,
+ 'b': (1, 2, 3)})
+ next(numbers)
+ self.assertEqual(inspect.getgeneratorlocals(numbers),
+ {'a': None, 'lst': [1, 2, 3], 'v': 3,
+ 'b': (1, 2, 3), 'c': 12})
+ try:
+ next(numbers)
+ except StopIteration:
+ pass
+ self.assertEqual(inspect.getgeneratorlocals(numbers), {})
+
+ def test_getgeneratorlocals_empty(self):
+ def yield_one():
+ yield 1
+ one = yield_one()
+ self.assertEqual(inspect.getgeneratorlocals(one), {})
+ try:
+ next(one)
+ except StopIteration:
+ pass
+ self.assertEqual(inspect.getgeneratorlocals(one), {})
+
+ def test_getgeneratorlocals_error(self):
+ self.assertRaises(TypeError, inspect.getgeneratorlocals, 1)
+ self.assertRaises(TypeError, inspect.getgeneratorlocals, lambda x: True)
+ self.assertRaises(TypeError, inspect.getgeneratorlocals, set)
+ self.assertRaises(TypeError, inspect.getgeneratorlocals, (2,3))
+
class TestSignatureObject(unittest.TestCase):
@staticmethod