bpo-45046: Support context managers in unittest (GH-28045)

Add methods enterContext() and enterClassContext() in TestCase.
Add method enterAsyncContext() in IsolatedAsyncioTestCase.
Add function enterModuleContext().
This commit is contained in:
Serhiy Storchaka 2022-05-08 17:49:09 +03:00 committed by GitHub
parent 8f29318079
commit 086c6b1b0f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
26 changed files with 307 additions and 92 deletions

View file

@ -102,12 +102,31 @@ def _id(obj):
return obj
def _enter_context(cm, addcleanup):
# We look up the special methods on the type to match the with
# statement.
cls = type(cm)
try:
enter = cls.__enter__
exit = cls.__exit__
except AttributeError:
raise TypeError(f"'{cls.__module__}.{cls.__qualname__}' object does "
f"not support the context manager protocol") from None
result = enter(cm)
addcleanup(exit, cm, None, None, None)
return result
_module_cleanups = []
def addModuleCleanup(function, /, *args, **kwargs):
"""Same as addCleanup, except the cleanup items are called even if
setUpModule fails (unlike tearDownModule)."""
_module_cleanups.append((function, args, kwargs))
def enterModuleContext(cm):
"""Same as enterContext, but module-wide."""
return _enter_context(cm, addModuleCleanup)
def doModuleCleanups():
"""Execute all module cleanup functions. Normally called for you after
@ -426,12 +445,25 @@ class TestCase(object):
Cleanup items are called even if setUp fails (unlike tearDown)."""
self._cleanups.append((function, args, kwargs))
def enterContext(self, cm):
"""Enters the supplied context manager.
If successful, also adds its __exit__ method as a cleanup
function and returns the result of the __enter__ method.
"""
return _enter_context(cm, self.addCleanup)
@classmethod
def addClassCleanup(cls, function, /, *args, **kwargs):
"""Same as addCleanup, except the cleanup items are called even if
setUpClass fails (unlike tearDownClass)."""
cls._class_cleanups.append((function, args, kwargs))
@classmethod
def enterClassContext(cls, cm):
"""Same as enterContext, but class-wide."""
return _enter_context(cm, cls.addClassCleanup)
def setUp(self):
"Hook method for setting up the test fixture before exercising it."
pass