How to use assertRaises in table-driven tests where some tests raise and others do not

Question

How can I avoid calling the function I'm testing in two different places when writing table-driven tests where some of the tests should raise but others should not?

This is what I want to do but it fails passing None to assertRaises:

tests = [
   (0, None),
   (1, None),
   (-1, TooFewException),
   (99, None),
   (100, TooManyException),
]
for n, exc in tests:
    with self.assertRaises(exc):
        results = my_code(n)
        assert len(results) == n

The best I have come up with is this but the redundant call to my_code is bothering me:

tests = [
   (0, None),
   (1, None),
   (-1, TooFewException),
   (99, None),
   (100, TooManyException),
]
for n, exc in tests:
    if exc is not None:
        with self.assertRaises(exc):
            my_code(n)
    else:
        results = my_code(n)
        assert len(results) == n

---

After adding a helper func on our base test case using the answer from @AmuroRay this is now:

tests = [
   (0, None),
   (1, None),
   (-1, TooFewException),
   (99, None),
   (100, TooManyException),
]
for n, exc in tests:
    with self.assertRaisesUnlessNone(exc):
        results = my_code(n)
        assert len(results) == n

Answer 1

contextlib has nullcontext() which can accomplish this. It's not really using fewer lines (without an ugly one-liner branch), but it eliminates the second user-code call:

from contextlib import nullcontext

for n, exc in tests:
    if exc is None:
        cm = nullcontext()
    else:
        cm = self.assertRaises(exc)
    with cm:
        results = my_code(n)
        assert len(results) == n