[TIP] assertRaisesRegex with exceptions from inside a generator

[Andrew Dalke]

Hi all,

  I tracked down an odd unit-test failure as part of my migration from Python 2.7 to 3.5+. The unittest assertRaisesRegex() does a traceback.clear_frames() on the traceback. If the exception comes from inside a still active generator then it has the effect of shutting down the generator.

I can't tell if this is misunderstanding on my part in how things are supposed to work, a bug in the unittest module, or some other bug in Python.

Before filing a Python bug, I want to eliminate the possibility of my misunderstanding. Note that all of the following was tested under Python 3.5.


My code uses a generator as a consumer coroutine, where I send data to the generator for output processing. Here's a sketch of how the generator looks:

  def generator():
    result = None
    while 1:
        process_list = yield result
        for item in process_list:
            process(item)
        result = ... intermediate result, like the number of records processed ...

(This is an unusual approach. The actual code is more complex with a lot of cumulative state information which must persist across multiple calls. There was a non-trivial performance penalty for saving/restoring through instance variables instead of using the STORE_FAST/LOAD_FAST by being in a generator's local scope.)

If there is an error then I want to raise an exception. I also want the consumer to be resumable, so the user can send more data to that generator.

Here's what the working code looks like:

##### Code under test

class MyException(Exception):
    pass

# The public interface is "open_writer()", which returns a Writer().
def open_writer():
    return Writer(write_gen())

class Writer:
    def __init__(self, gen):
        self.gen = gen
    def write_objects(self, objs):
        total_len, err = self.gen.send(objs)
        # If the generator indicates there was an error then re-raise it
        if err is not None:
            raise err
        return total_len

# The generator is a back-end implementation detail.

def write_gen():
    gen = _write_gen()
    next(gen) # prime the generator so it's ready to receive
    return gen

def _write_gen():
    total_len = 0 # accumulate the length of all of the objects
    result = "started"

    # Receive a request to process more objects.
    while 1:
        objs = yield result
        try:
            for obj in objs:
                if obj == "a":
                    raise MyException("'a' is not allowed")

                total_len += len(obj)

        except Exception as err:
            # An error occurred. Forward it to the Writer() instance.
            result = None, err
        else:
            # No error. Return the current total length.
            result = total_len, None
##### End of code under test

Here's a simple test for the code:

##### test the code manually. This works as expected
def simple_test():
    w = open_writer()
    try:
        w.write_objects(["this", "is", "a"])
    except MyException as err:
        print("got expected exception:", err)
    else:
        raise AssertionError("should have raised exception")

    total_len = w.write_objects(["test"])
    assert total_len == 10, total_len
    print("simple_test is done.")

simple_test()
##### end of manual test code

The output of this is exactly as expected.

##### Output of the manual test code.
got expected exception: 'a' is not allowed
simple_test is done.
#####


My problem came because the real test code is in unittest and I use the assertRaisesRegex to test the exception. That causes the generator to stop working, and not be resumable.


##### test the code using unittest

import unittest
class SimpleText(unittest.TestCase):
    def test_using_assert_raises(self):
        w = open_writer()
        with self.assertRaisesRegex(MyException, "'a' is not allowed"):
            w.write_objects(["this", "is", "a"])

        total_len = w.write_objects(["test"])
        self.assertEqual(total_len, 10)

try:
    unittest.main()
except SystemExit:
    pass
##### End of the unittest-based test.

This gives me an unexpected StopIteration error.

##### Output from the unittest-based test:
E
======================================================================
ERROR: test_using_assert_raises (__main__.SimpleText)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "fail2.py", line 69, in test_using_assert_raises
    total_len = w.write_objects(["test"])
  File "fail2.py", line 13, in write_objects
    total_len, err = self.gen.send(objs)
StopIteration

----------------------------------------------------------------------
Ran 1 test in 0.001s

FAILED (errors=1)
##### End of the output from unittest.

I spent about 5-6 hours trying to figure out if I didn't understand Python 3's chained exceptions. Eventually I tracked it down to unittest.py::_AssertRaisesContext.__exit__'s use of traceback.clear_frames(tb), at

  https://hg.python.org/cpython/file/tip/Lib/unittest/case.py#l201

Here's my reproducible which shows how that function causes the behavior I see:

##### Test how clear_frames() affects a traceback from an active generator.
import traceback

class ClearFrames:
    def __init__(self, clear_frames):
        self.clear_frames = clear_frames
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        assert exc_type is MyException, exc_type
        assert "is not allowed" in str(exc_value), str(exc_value)
        if self.clear_frames:
            traceback.clear_frames(tb)  # This is the only difference between the tests.
        return True

# This is essentially the same test case as before, but structured using
# a context manager that either does or does not clear the traceback frames.
def clear_test(clear_frames):
    w = open_writer()
    with ClearFrames(clear_frames):
        w.write_objects(["this", "is", "a"])
    try:
        total_len = w.write_objects(["test"])
    except StopIteration:
        print("...caught StopIteration")
    else:
        assert total_len == 10, total_len
        print("...returned a length")

print("\nDo not clear frames")
clear_test(False)
print("\nClear frames")
clear_test(True)
##### End of ClearFrames() context manager test code.

This test code results in:

##### Output from the ClearFrames() context manager test:
Do not clear frames
...returned a length

Clear frames
...caught StopIteration
##### End of ClearFrames() test output.



It looks like traceback.clear_frames(tb) is the culprit. ("Clears the local variables of all the stack frames in a traceback tb by calling the clear() method of each frame object.") I assume it expects that the frame will never be used again, which isn't the case for an exception from a generator.

My work-around is to not use assertRaisesRegex for my unit tests for resumable consumer generators. I only have a dozen or so of these so it's easy to re-write.

Is there a better solution? I can think of several:

  - is there a way I can sanitize the exception in my code so users can see the traceback stack but not have clear_frames() destroy things?

  - should unittest2 not clear_frames() for a frame from a running generator?

  - should assertRaisesRegex take an optional parameter to not clear the frames?

  - should traceback.clear_frames() only clear up to a running generator?

  - Perhaps there's a missing function which would be more appropriate than clear_frames() for this use case?

Cheers,


				Andrew
				dalke at dalkescientific.com