Skip to content

Commit

Permalink
Extract decorator logic into separate methods
Browse files Browse the repository at this point in the history
So that it's easier to re-use callable decoration as part of the
updated class decorator.
  • Loading branch information
leokon committed Sep 28, 2021
1 parent df41aec commit e8ab095
Showing 1 changed file with 50 additions and 45 deletions.
95 changes: 50 additions & 45 deletions src/time_machine/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,56 +291,61 @@ def __call__(
Callable[..., Any],
]:
if isinstance(wrapped, type):
# Class decorator
if not issubclass(wrapped, TestCase):
raise TypeError("Can only decorate unittest.TestCase subclasses.")

# Modify the setUpClass method
orig_setUpClass = wrapped.setUpClass

@functools.wraps(orig_setUpClass)
def setUpClass(cls: Type[TestCase]) -> None:
self.__enter__()
try:
orig_setUpClass()
except Exception:
self.__exit__(*sys.exc_info())
raise

wrapped.setUpClass = classmethod(setUpClass) # type: ignore[assignment]

orig_tearDownClass = wrapped.tearDownClass

@functools.wraps(orig_tearDownClass)
def tearDownClass(cls: Type[TestCase]) -> None:
orig_tearDownClass()
self.__exit__(None, None, None)

wrapped.tearDownClass = classmethod( # type: ignore[assignment]
tearDownClass
)
return wrapped
return self._decorate_class(wrapped)
elif inspect.iscoroutinefunction(wrapped):

@functools.wraps(wrapped)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
with self:
# mypy has not narrowed 'wrapped' to a coroutine function
return await wrapped(
*args,
**kwargs,
) # type: ignore [misc,operator]

return wrapper
return self._decorate_coroutine(wrapped)
else:
assert callable(wrapped)
return self._decorate_callable(wrapped)

def _decorate_class(self, wrapped: Type[TestCase]) -> Type[TestCase]:
if not issubclass(wrapped, TestCase):
raise TypeError("Can only decorate unittest.TestCase subclasses.")

# Modify the setUpClass method
orig_setUpClass = wrapped.setUpClass

@functools.wraps(orig_setUpClass)
def setUpClass(cls: Type[TestCase]) -> None:
self.__enter__()
try:
orig_setUpClass()
except Exception:
self.__exit__(*sys.exc_info())
raise

@functools.wraps(wrapped)
def wrapper(*args: Any, **kwargs: Any) -> Any:
with self:
return wrapped(*args, **kwargs)
wrapped.setUpClass = classmethod(setUpClass) # type: ignore[assignment]

return wrapper
orig_tearDownClass = wrapped.tearDownClass

@functools.wraps(orig_tearDownClass)
def tearDownClass(cls: Type[TestCase]) -> None:
orig_tearDownClass()
self.__exit__(None, None, None)

wrapped.tearDownClass = classmethod(tearDownClass) # type: ignore[assignment]
return wrapped

def _decorate_coroutine(
self, wrapped: Callable[..., Coroutine[Any, Any, Any]]
) -> Callable[..., Coroutine[Any, Any, Any]]:
@functools.wraps(wrapped)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
with self:
return await wrapped(
*args,
**kwargs,
)

return wrapper

def _decorate_callable(self, wrapped: Callable[..., Any]) -> Callable[..., Any]:
@functools.wraps(wrapped)
def wrapper(*args: Any, **kwargs: Any) -> Any:
with self:
return wrapped(*args, **kwargs)

return wrapper


# datetime module
Expand Down

0 comments on commit e8ab095

Please sign in to comment.