Support non-unittest class decoration

Wraps all class methods that begin with "test", matching default
pytest method collection.

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,15 +281,11 @@ 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):
             return self._decorate_class(wrapped)
         elif inspect.iscoroutinefunction(wrapped):
@@ -298,32 +294,44 @@ def __call__(
             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.")
+    def _decorate_class(self, wrapped: type) -> type:
+        if issubclass(wrapped, TestCase):
+            # Modify the setUpClass method
+            orig_setUpClass = wrapped.setUpClass
 
-        # 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(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]
 
-        wrapped.setUpClass = classmethod(setUpClass)  # type: ignore[assignment]
+            orig_tearDownClass = wrapped.tearDownClass
 
-        orig_tearDownClass = wrapped.tearDownClass
+            @functools.wraps(orig_tearDownClass)
+            def tearDownClass(cls: Type[TestCase]) -> None:
+                orig_tearDownClass()
+                self.__exit__(None, None, None)
 
-        @functools.wraps(orig_tearDownClass)
-        def tearDownClass(cls: Type[TestCase]) -> None:
-            orig_tearDownClass()
-            self.__exit__(None, None, None)
+            wrapped.tearDownClass = classmethod(  # type: ignore[assignment]
+                tearDownClass
+            )
+        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
+
+                if not callable(value) or inspect.isclass(value):
+                    continue
+
+                wrapped_method = self._decorate_callable(value)
+                setattr(wrapped, attribute, wrapped_method)
 
-        wrapped.tearDownClass = classmethod(tearDownClass)  # type: ignore[assignment]
         return wrapped
 
     def _decorate_coroutine(

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: