Add support for generators as pytest fixture classes

README.md

@@ -184,6 +184,88 @@ Fixture classes fully support async `__call__` methods. Simply define `__call__`
 
 Note: Async fixtures require `pytest-asyncio` or another async pytest plugin to be installed.
 
+### Lifespan Support
+
+Fixture classes support lifespan methods for setup and teardown logic, similar to pytest's built-in fixture teardown. This is useful when your factory needs to perform initialization before it can be used and cleanup after the test completes.
+
+#### Synchronous Lifespan
+
+Define a `lifespan` method that returns a `Generator` to add setup and teardown logic to your fixture class:
+
+```python
+from pytest_fixture_classes import fixture_class
+from collections.abc import Generator
+import pytest
+
+@pytest.fixture
+def database_url() -> str:
+    return "sqlite:///test.db"
+
+@fixture_class(name="user_factory")
+class UserFactory:
+    database_url: str
+
+    def lifespan(self) -> Generator:
+        # Setup: runs when the fixture is first created
+        print(f"Connecting to {self.database_url}")
+        self.connection = create_connection(self.database_url)
+
+        yield  # Test runs here
+
+        # Teardown: runs after the test completes
+        print("Closing database connection")
+        self.connection.close()
+
+    def __call__(self, name: str, email: str) -> User:
+        return User.create(self.connection, name=name, email=email)
+
+def test_users(user_factory: UserFactory):
+    user1 = user_factory("Alice", "[email protected]")
+    user2 = user_factory("Bob", "[email protected]")
+    # Connection will be closed after this test completes
+```
+
+#### Asynchronous Lifespan
+
+For async fixture classes, define an `async def lifespan` method that returns an `AsyncGenerator`:
+
+```python
+from pytest_fixture_classes import fixture_class
+from collections.abc import AsyncGenerator
+import pytest_asyncio
+import pytest
+
+@pytest_asyncio.fixture
+async def async_database_url() -> str:
+    return "sqlite:///test.db"
+
+@fixture_class(name="async_user_factory")
+class AsyncUserFactory:
+    async_database_url: str
+
+    async def lifespan(self) -> AsyncGenerator:
+        # Async setup
+        print(f"Connecting to {self.async_database_url}")
+        self.connection = await create_async_connection(self.async_database_url)
+
+        yield  # Test runs here
+
+        # Async teardown
+        print("Closing async database connection")
+        await self.connection.close()
+
+    async def __call__(self, name: str, email: str) -> User:
+        return await User.create(self.connection, name=name, email=email)
+
+@pytest.mark.asyncio
+async def test_async_users(async_user_factory: AsyncUserFactory):
+    user1 = await async_user_factory("Alice", "[email protected]")
+    user2 = await async_user_factory("Bob", "[email protected]")
+    # Connection will be closed after this test completes
+```
+
+The lifespan method is optional. If you don't need setup/teardown logic, you can omit it entirely.
+
 ## Implementation details
 
 * The fixture_class decorator turns your class into a frozen dataclass with slots so you won't be able to add new attributes to it after definiton. You can, however, define any methods you like except `__init__`.

pytest_fixture_classes/fixture_class.py

@@ -7,6 +7,13 @@
 import pytest
 from typing_extensions import Literal, dataclass_transform
 
+try:
+    import pytest_asyncio
+
+    HAS_PYTEST_ASYNCIO = True
+except ImportError:
+    HAS_PYTEST_ASYNCIO = False
+
 T = TypeVar("T", bound=type)
 
 _ScopeName = Literal["session", "package", "module", "class", "function"]
@@ -58,19 +65,67 @@ def inner(fixture_cls):
             fixture_cls
         )
         args = list(inspect.signature(fixture_dataclass.__init__).parameters)[1:]
-        func_def = dedent(
-            f"""
-            def {fixture_dataclass.__name__}({', '.join(args)}):
-                return fixture_cls({', '.join(args)})
-        """
-        )
+
+        has_lifespan = hasattr(fixture_cls, "lifespan")
+        is_async_lifespan = False
+
+        if has_lifespan:
+            lifespan_method = getattr(fixture_cls, "lifespan")
+            is_async_lifespan = inspect.isasyncgenfunction(lifespan_method)
+
+            if is_async_lifespan:
+                func_def = dedent(
+                    f"""
+                    async def {fixture_dataclass.__name__}({', '.join(args)}):
+                        instance = fixture_cls({', '.join(args)})
+                        lifespan_gen = instance.lifespan()
+                        await lifespan_gen.__anext__()
+                        yield instance
+                        try:
+                            await lifespan_gen.__anext__()
+                        except StopAsyncIteration:
+                            pass
+                    """
+                )
+            else:
+                func_def = dedent(
+                    f"""
+                    def {fixture_dataclass.__name__}({', '.join(args)}):
+                        instance = fixture_cls({', '.join(args)})
+                        lifespan_gen = instance.lifespan()
+                        next(lifespan_gen)
+                        yield instance
+                        try:
+                            next(lifespan_gen)
+                        except StopIteration:
+                            pass
+                    """
+                )
+        else:
+            func_def = dedent(
+                f"""
+                def {fixture_dataclass.__name__}({', '.join(args)}):
+                    return fixture_cls({', '.join(args)})
+                """
+            )
+
         namespace = {"fixture_cls": fixture_dataclass}
         exec(func_def, namespace)
         func = namespace[fixture_dataclass.__name__]
         func.__module__ = fixture_cls.__module__
         func.__doc__ = fixture_cls.__doc__ or fixture_dataclass.__doc__
-        return pytest.fixture(
-            scope=scope, params=params, autouse=autouse, ids=ids, name=name
-        )(func)
+
+        if has_lifespan and is_async_lifespan:
+            if not HAS_PYTEST_ASYNCIO:
+                raise ImportError(
+                    "pytest-asyncio is required for async lifespan support"
+                )
+            return pytest_asyncio.fixture(
+                scope=scope, params=params, autouse=autouse, ids=ids, name=name
+            )(func)
+        else:
+            return pytest.fixture(
+                scope=scope, params=params, autouse=autouse, ids=ids, name=name
+            )(func)
 
     return inner if fixture_cls is None else inner(fixture_cls)

tests/test_async_fixture_class.py

@@ -1,4 +1,5 @@
 import asyncio
+
 import pytest
 
 from pytest_fixture_classes import fixture_class
@@ -11,7 +12,6 @@ def sync_fixture() -> int:
 
 @fixture_class(name="async_factory")
 class AsyncFactory:
-    """Async factory fixture"""
     sync_fixture: int
 
     async def __call__(self, multiplier: int) -> int:
@@ -43,5 +43,7 @@ async def test_nested_async_factory(nested_async_factory: NestedAsyncFactory) ->
 
 
 @pytest.mark.asyncio
-async def test_async_factory_docstring_preservation(async_factory: AsyncFactory) -> None:
+async def test_async_factory_docstring_preservation(
+    async_factory: AsyncFactory,
+) -> None:
     assert async_factory.__doc__ == "Async factory fixture"

tests/test_lifespan_fixture_class.py

@@ -0,0 +1,75 @@
+import asyncio
+from collections.abc import AsyncGenerator, Generator
+
+import pytest
+import pytest_asyncio
+
+from pytest_fixture_classes import fixture_class
+
+execution_log: list[str] = []
+
+
+@pytest.fixture(autouse=True)
+def reset_execution_log():
+    execution_log.clear()
+
+
+@pytest.fixture
+def sync_fixture() -> int:
+    return 100
+
+
+@pytest_asyncio.fixture
+async def async_fixture() -> int:
+    await asyncio.sleep(0)
+    return 2
+
+
+@fixture_class(name="sync_lifespan_factory")
+class SyncLifespanFactory:
+    sync_fixture: int
+
+    def lifespan(self) -> Generator:
+        execution_log.append("sync_factory_init")
+        yield
+        execution_log.append("sync_factory_cleanup")
+
+    def __call__(self, value: int) -> int:
+        return value * self.sync_fixture
+
+
+@fixture_class(name="async_lifespan_factory")
+class AsyncLifespanFactory:
+    sync_fixture: int
+    async_fixture: int
+
+    async def lifespan(self) -> AsyncGenerator:
+        execution_log.append("async_factory_init")
+        yield
+        execution_log.append("async_factory_cleanup")
+
+    async def __call__(self, value: int) -> int:
+        return value * self.sync_fixture + self.async_fixture
+
+
+def test_sync_lifespan_factory(sync_lifespan_factory: SyncLifespanFactory) -> None:
+    assert "sync_factory_init" in execution_log
+
+    assert sync_lifespan_factory(1) == 100
+    assert sync_lifespan_factory(2) == 200
+    assert sync_lifespan_factory(3) == 300
+
+    assert "sync_factory_cleanup" not in execution_log
+
+
+@pytest.mark.asyncio
+async def test_async_lifespan_factory(
+    async_lifespan_factory: AsyncLifespanFactory,
+) -> None:
+    assert "async_factory_init" in execution_log
+
+    assert await async_lifespan_factory(1) == 102
+    assert await async_lifespan_factory(2) == 202
+    assert await async_lifespan_factory(3) == 302
+
+    assert "async_factory_cleanup" not in execution_log