Convert cleanups to generators

README.md

@@ -12,7 +12,7 @@
 > At this point, it's unclear whether this project will become a "proper Hynek project".
 > I will keep using it for my work projects, but whether this will grow beyond my personal needs depends on community interest.
 
-*svc-reg* is a [service locator](https://en.wikipedia.org/wiki/Service_locator_pattern) for Python that lets you register factories for types/interfaces and then create instances of those types with unified life-cycle management and health checks.
+*svc-reg* is a [service locator](https://en.wikipedia.org/wiki/Service_locator_pattern) for Python that lets you register factories for types/interfaces and then create instances of those types with **unified life-cycle management** and **health checks**.
 
 **This allows you to configure and manage resources in *one central place* and access them in a *consistent* way.**
 
@@ -48,6 +48,20 @@ def view():
 
 The latter already works with [Flask](#flask).
 
+And you set it up like this:
+
+```python
+engine = create_engine("postgresql://localhost")
+
+def engine_factory():
+    with engine.connect() as conn:
+        yield conn
+
+svc_reg.register_factory(Database, create_database)
+```
+
+The generator-based setup and cleanup may remind you of [Pytest fixtures](https://docs.pytest.org/en/stable/explanation/fixtures.html).
+
 <!-- end-pypi -->
 
 
@@ -97,11 +111,13 @@ True
 ```
 
 A container lives as long as you want the instances to live -- e.g., as long as a request lives.
-At the end, you run `container.close()` to clean up all instances that the container has created.
+
+If a factory is a generator and yields the instance, the generator will be remembered.
+At the end, you run `container.close()` and all generators will be finished (i.e. called `next(g)`).
 You can use this to return database connections to a pool, et cetera.
 
-If you have async cleanup functions, use `await container.aclose()` instead.
-It will run both sync and async cleanup functions.
+If you have async generators, use `await container.aclose()` instead which calls `await g.__anext__()` on all async generators.
+It will run both sync and async cleanup functions by default.
 
 Failing cleanups are logged at `warning` level but otherwise ignored.
 
@@ -162,24 +178,29 @@ def create_app(config_filename):
     app = svc_reg.flask.init_app(app)
 
     # Now, register a factory that calls `engine.connect()` if you ask for a
-    # Connections and `connection.close()` on cleanup.
+    # Connections. Since we use yield inside of a context manager, the
+    # connection gets cleaned up when the container is closed.
     # If you ask for a ping, it will run `SELECT 1` on a new connection and
     # clean up the connection behind itself.
     engine = create_engine("postgresql://localhost")
+    def engine_factory():
+        with engine.connect() as conn:
+            yield conn
+
     ping = text("SELECT 1")
     svc_reg_flask.register_factory(
         # The app argument makes it good for custom init_app() functions.
         app,
         Connection,
-        engine.connect,
-        cleanup=lambda conn: conn.close(),
+        engine_factory,
         ping=lambda conn: conn.execute(ping)
     )
 
     # You also use svc_reg WITHIN factories:
     svc_reg_flask.register_factory(
         app, # <---
         AbstractRepository,
+        # No cleanup, so we just return an object using a lambda
         lambda: Repository.from_connection(
             svc_reg.flask.get(Connection)
         ),

conftest.py

@@ -1,8 +1,12 @@
 from doctest import ELLIPSIS
 
+import pytest
+
 from sybil import Sybil
 from sybil.parsers.rest import DocTestParser, PythonCodeBlockParser
 
+import svc_reg
+
 
 pytest_collect_file = Sybil(
     parsers=[
@@ -11,3 +15,13 @@
     ],
     patterns=["*.md", "*.py"],
 ).pytest()
+
+
+@pytest.fixture(name="registry")
+def _registry():
+    return svc_reg.Registry()
+
+
+@pytest.fixture(name="container")
+def _container(registry):
+    return svc_reg.Container(registry)

pyproject.toml

@@ -157,6 +157,10 @@ ignore = [
   "EM101",  # no need for exception msg hygiene in tests
 ]
 
+"conftest.py" = [
+  "PT005", # we always add underscores and explicit name
+]
+
 
 [tool.ruff.isort]
 lines-between-types = 1

src/svc_reg/__init__.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+from . import exceptions
 from ._core import (
     Container,
     RegisteredService,
     Registry,
-    ServiceNotFoundError,
     ServicePing,
 )
 
@@ -15,6 +15,7 @@
     "Registry",
     "ServiceNotFoundError",
     "ServicePing",
+    "exceptions",
 ]
 
 try:

src/svc_reg/_core.py

@@ -1,22 +1,18 @@
 from __future__ import annotations
 
-import asyncio
 import logging
+import warnings
 
 from collections.abc import Callable
 from contextlib import suppress
-from typing import Any
+from typing import Any, AsyncGenerator, Generator
 
 import attrs
 
-
-log = logging.getLogger(__name__)
+from .exceptions import ServiceNotFoundError
 
 
-class ServiceNotFoundError(Exception):
-    """
-    Raised when a service type is not registered.
-    """
+log = logging.getLogger(__name__)
 
 
 @attrs.define
@@ -27,16 +23,16 @@ class Container:
 
     registry: Registry
     instantiated: dict[type, object] = attrs.Factory(dict)
-    cleanups: list[tuple[RegisteredService, object]] = attrs.Factory(list)
-    async_cleanups: list[tuple[RegisteredService, object]] = attrs.Factory(
-        list
-    )
+    cleanups: list[tuple[RegisteredService, Generator]] = attrs.Factory(list)
+    async_cleanups: list[
+        tuple[RegisteredService, AsyncGenerator]
+    ] = attrs.Factory(list)
 
     def __repr__(self) -> str:
         return (
             f"<Container(instantiated={len(self.instantiated)}, "
             f"cleanups={len(self.cleanups)}, "
-            f"async_cleanups={len(self.async_cleanups)}>"
+            f"async_cleanups={len(self.async_cleanups)})>"
         )
 
     def get(self, svc_type: type) -> Any:
@@ -48,42 +44,44 @@ def get(self, svc_type: type) -> Any:
         Returns:
              Any until https://github.com/python/mypy/issues/4717 is fixed.
         """
-        if (svc := self._get_instance(svc_type)) is not None:
+        if (svc := self.instantiated.get(svc_type)) is not None:
             return svc
 
         rs = self.registry.get_registered_service_for(svc_type)
         svc = rs.factory()
-        self._add_instance(rs, svc)
 
-        return svc
+        if isinstance(svc, Generator):
+            self.cleanups.append((rs, svc))
+            svc = next(svc)
 
-    def _add_instance(self, rs: RegisteredService, svc: object) -> None:
         self.instantiated[rs.svc_type] = svc
-        self.add_cleanup(rs, svc)
 
-    def _get_instance(self, svc_type: type) -> object | None:
-        """
-        If present, return the cached instance of *svc_type*.
-        """
-        return self.instantiated.get(svc_type)
+        return svc
 
-    def add_cleanup(self, rs: RegisteredService, svc: object) -> bool:
+    async def aget(self, svc_type: type) -> Any:
         """
-        Add a cleanup function for *svc* if *rs* has one without remembering
-        the service itself.
+        Get an instance of *svc_type*.
+
+        Instantiate it asynchronously if necessary and register its cleanup.
 
-        Return:
-            True if a cleanup was added, False otherwise.
+        Returns:
+             Any until https://github.com/python/mypy/issues/4717 is fixed.
         """
-        if not rs.cleanup:
-            return False
+        if (svc := self.instantiated.get(svc_type)) is not None:
+            return svc
 
-        if asyncio.iscoroutinefunction(rs.cleanup):
+        rs = self.registry.get_registered_service_for(svc_type)
+        svc = rs.factory()
+
+        if isinstance(svc, AsyncGenerator):
             self.async_cleanups.append((rs, svc))
+            svc = await svc.__anext__()
         else:
-            self.cleanups.append((rs, svc))
+            svc = await svc  # type: ignore[misc]
 
-        return True
+        self.instantiated[rs.svc_type] = svc
+
+        return svc
 
     def forget_service_type(self, svc_type: type) -> None:
         """
@@ -97,10 +95,16 @@ def close(self) -> None:
         Run all synchronous registered cleanups.
         """
         while self.cleanups:
-            rs, svc = self.cleanups.pop()
+            rs, gen = self.cleanups.pop()
             try:
-                rs.cleanup(svc)  # type: ignore[misc]
-            except Exception:  # noqa: PERF203, BLE001
+                next(gen)
+
+                warnings.warn(
+                    f"clean up for {rs!r} didn't stop iterating", stacklevel=1
+                )
+            except StopIteration:  # noqa: PERF203
+                pass
+            except Exception:  # noqa: BLE001
                 log.warning(
                     "clean up failed",
                     exc_info=True,
@@ -114,10 +118,17 @@ async def aclose(self) -> None:
         self.close()
 
         while self.async_cleanups:
-            rs, svc = self.async_cleanups.pop()
+            rs, gen = self.async_cleanups.pop()
             try:
-                await rs.cleanup(svc)  # type: ignore[misc]
-            except Exception:  # noqa: PERF203, BLE001
+                await gen.__anext__()
+
+                warnings.warn(
+                    f"clean up for {rs!r} didn't stop iterating", stacklevel=1
+                )
+
+            except StopAsyncIteration:  # noqa: PERF203
+                pass
+            except Exception:  # noqa: BLE001
                 log.warning(
                     "clean up failed",
                     exc_info=True,
@@ -139,7 +150,6 @@ def get_pings(self) -> list[ServicePing]:
 class RegisteredService:
     svc_type: type
     factory: Callable = attrs.field(hash=False)
-    cleanup: Callable | None = attrs.field(hash=False)
     ping: Callable | None = attrs.field(hash=False)
 
     @property
@@ -150,7 +160,6 @@ def __repr__(self) -> str:
         return (
             f"<RegisteredService(svc_type={ self.svc_type.__module__ }."
             f"{ self.svc_type.__qualname__ }, "
-            f"has_cleanup={ self.cleanup is not None}, "
             f"has_ping={ self.ping is not None})>"
         )
 
@@ -161,8 +170,7 @@ class ServicePing:
     _rs: RegisteredService
 
     def ping(self) -> None:
-        svc = self._rs.factory()
-        self._container.add_cleanup(self._rs, svc)
+        svc = self._container.get(self._rs.svc_type)
         self._rs.ping(svc)  # type: ignore[misc]
 
     @property
@@ -179,24 +187,18 @@ def register_factory(
         svc_type: type,
         factory: Callable,
         *,
-        cleanup: Callable | None = None,
         ping: Callable | None = None,
     ) -> None:
-        self.services[svc_type] = RegisteredService(
-            svc_type, factory, cleanup, ping
-        )
+        self.services[svc_type] = RegisteredService(svc_type, factory, ping)
 
     def register_value(
         self,
         svc_type: type,
         instance: object,
         *,
-        cleanup: Callable | None = None,
         ping: Callable | None = None,
     ) -> None:
-        self.register_factory(
-            svc_type, lambda: instance, cleanup=cleanup, ping=ping
-        )
+        self.register_factory(svc_type, lambda: instance, ping=ping)
 
     def get_registered_service_for(self, svc_type: type) -> RegisteredService:
         try:

src/svc_reg/exceptions.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+
+class ServiceNotFoundError(Exception):
+    """
+    Raised when a service type is not registered.
+    """