Support non-unittest class decoration

README.rst

@@ -178,8 +178,7 @@ When used as a context manager, time is mocked during the ``with`` block:
 Class Decorator
 ^^^^^^^^^^^^^^^
 
-Only ``unittest.TestCase`` subclasses are supported.
-When applied as a class decorator to such classes, time is mocked from the start of ``setUpClass()`` to the end of ``tearDownClass()``:
+Decorating a ``unittest.TestCase`` subclass will mock time from the start of ``setUpClass()`` to the end of ``tearDownClass()``:
 
 .. code-block:: python
 
@@ -195,6 +194,22 @@ When applied as a class decorator to such classes, time is mocked from the start
 
 Note this is different to ``unittest.mock.patch()``\'s behaviour, which is to mock only during the test methods.
 
+Any other class can also be decorated - time will be mocked for each method with a name starting with ``test``:
+
+.. code-block:: python
+
+    import time
+    import time_machine
+
+
+    @time_machine.travel(0.0)
+    class TestDeepPast:
+        def test_in_the_deep_past(self):
+            assert 0.0 < time.time() < 1.0
+
+        def not_a_test(self):
+            assert not 0.0 < time.time() < 1.0
+
 Timezone mocking
 ^^^^^^^^^^^^^^^^
 

src/time_machine/__init__.py

@@ -263,7 +263,7 @@ def __exit__(
         self.stop()
 
     @overload
-    def __call__(self, wrapped: Type[TestCase]) -> Type[TestCase]:  # pragma: no cover
+    def __call__(self, wrapped: type) -> type:  # pragma: no cover
         ...
 
     @overload
@@ -281,20 +281,21 @@ def __call__(
     def __call__(
         self,
         wrapped: Union[
-            Type[TestCase],
+            type,
             Callable[..., Coroutine[Any, Any, Any]],
             Callable[..., Any],
         ],
-    ) -> Union[
-        Type[TestCase],
-        Callable[..., Coroutine[Any, Any, Any]],
-        Callable[..., Any],
-    ]:
+    ) -> Union[type, Callable[..., Coroutine[Any, Any, Any]], Callable[..., Any]]:
         if isinstance(wrapped, type):
-            # Class decorator
-            if not issubclass(wrapped, TestCase):
-                raise TypeError("Can only decorate unittest.TestCase subclasses.")
+            return self._decorate_class(wrapped)
+        elif inspect.iscoroutinefunction(wrapped):
+            return self._decorate_coroutine(wrapped)
+        else:
+            assert callable(wrapped)
+            return self._decorate_callable(wrapped)
 
+    def _decorate_class(self, wrapped: type) -> type:
+        if issubclass(wrapped, TestCase):
             # Modify the setUpClass method
             orig_setUpClass = wrapped.setUpClass
 
@@ -319,28 +320,40 @@ def tearDownClass(cls: Type[TestCase]) -> None:
             wrapped.tearDownClass = classmethod(  # type: ignore[assignment]
                 tearDownClass
             )
-            return wrapped
-        elif inspect.iscoroutinefunction(wrapped):
+        else:
+            # Wrap all methods starting with "test", intended for pytest-like classes
+            for (attribute, value) in wrapped.__dict__.items():
+                if not attribute.startswith("test"):
+                    continue
 
-            @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]
+                if not callable(value) or inspect.isclass(value):
+                    continue
 
-            return wrapper
-        else:
-            assert callable(wrapped)
+                wrapped_method = self._decorate_callable(value)
+                setattr(wrapped, attribute, wrapped_method)
 
-            @functools.wraps(wrapped)
-            def wrapper(*args: Any, **kwargs: Any) -> Any:
-                with self:
-                    return wrapped(*args, **kwargs)
+        return wrapped
 
-            return wrapper
+    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

tests/test_time_machine.py

@@ -487,14 +487,35 @@ async def record_time():
     assert recorded_time == EPOCH + 140.0
 
 
-def test_class_decorator_fails_non_testcase():
-    with pytest.raises(TypeError) as excinfo:
+@time_machine.travel(EPOCH)
+def test_non_unittest_class_decorator():
+    @time_machine.travel(EPOCH + 25.0)
+    class NonUnitTestClassTests:
+        def test_class_decorator(self):
+            assert time.time() == EPOCH + 25.0
 
-        @time_machine.travel(EPOCH)
-        class Something:
-            pass
+    NonUnitTestClassTests().test_class_decorator()
+
+
+@time_machine.travel(EPOCH)
+def test_non_unittest_class_decorator_ignores_non_tests():
+    @time_machine.travel(EPOCH + 25.0)
+    class NonUnitTestClassTests:
+        def not_a_test(self):
+            assert time.time() == EPOCH
+
+    NonUnitTestClassTests().not_a_test()
+
+
+@time_machine.travel(EPOCH)
+def test_non_unittest_class_decorator_ignores_nested_classes():
+    @time_machine.travel(EPOCH + 25.0)
+    class NonUnitTestClassTests:
+        class test_class:
+            def test(self):
+                assert time.time() == EPOCH
 
-    assert excinfo.value.args == ("Can only decorate unittest.TestCase subclasses.",)
+    NonUnitTestClassTests().test_class().test()
 
 
 class MethodDecoratorTests: