Close #15153: Added inspect.getgeneratorlocals to simplify whitebox testing of generator state updates

5 files changed, 98 insertions, 0 deletions

diff --git a/Doc/library/inspect.rst b/Doc/library/inspect.rst
index 04b724b..6568e94 100644
--- a/Doc/library/inspect.rst
+++ b/Doc/library/inspect.rst
@@ -676,3 +676,27 @@ generator to be determined easily.
     * GEN_CLOSED: Execution has completed.
 
    .. versionadded:: 3.2
+
+The current internal state of the generator can also be queried. This is
+mostly useful for testing purposes, to ensure that internal state is being
+updated as expected:
+
+.. function:: getgeneratorlocals(generator)
+
+   Get the mapping of live local variables in *generator* to their current
+   values.  A dictionary is returned that maps from variable names to values.
+   This is the equivalent of calling :func:`locals` in the body of the
+   generator, and all the same caveats apply.
+
+   If *generator* is a :term:`generator` with no currently associated frame,
+   then an empty dictionary is returned.  :exc:`TypeError` is raised if
+   *generator* is not a Python generator object.
+
+   .. impl-detail::
+
+      This function relies on the generator exposing a Python stack frame
+      for introspection, which isn't guaranteed to be the case in all
+      implementations of Python. In such cases, this function will always
+      return an empty dictionary.
+
+   .. versionadded:: 3.3
diff --git a/Doc/whatsnew/3.3.rst b/Doc/whatsnew/3.3.rst
index 592f9c9..c6225c3 100644
--- a/Doc/whatsnew/3.3.rst
+++ b/Doc/whatsnew/3.3.rst
@@ -1037,6 +1037,13 @@ state when testing code that relies on stateful closures.
 
 (Contributed by Meador Inge and Nick Coghlan in :issue:`13062`)
 
+A new :func:`~inspect.getgeneratorlocals` function has been added. This
+function reports the current binding of local variables in the generator's
+stack frame, making it easier to verify correct internal state when testing
+generators.
+
+(Contributed by Meador Inge in :issue:`15153`)
+
 io
 --
 
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 @@ def getattr_static(obj, attr, default=_sentinel):
     raise AttributeError(attr)
 
 
+# ------------------------------------------------ generator introspection
+
 GEN_CREATED = 'GEN_CREATED'
 GEN_RUNNING = 'GEN_RUNNING'
 GEN_SUSPENDED = 'GEN_SUSPENDED'
@@ -1282,6 +1284,22 @@ def getgeneratorstate(generator):
     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 @@ class TestGetGeneratorState(unittest.TestCase):
             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
diff --git a/Misc/NEWS b/Misc/NEWS
index 0c69f35..37a0dee 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -40,6 +40,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #15153: Added inspect.getgeneratorlocals to simplify white box
+  testing of generator state updates
+
 - Issue #13062: Added inspect.getclosurevars to simplify testing stateful
   closures