Add lifespan helper (#50)

* Add `lifespan` helper

CHANGELOG.md

@@ -6,6 +6,10 @@ We follow [Semantic Versions](https://semver.org/).
 
 -
 
+## Version 0.14.0
+
+- Added `helpers.lifespan` function for simple resource management
+
 ## Version 0.13.0
 
 - Fix async singleton through sync resolving

README.md

@@ -707,6 +707,65 @@ def get_connection(
     print("connecting to", host, port)
 ```
 
+#### `helpers.lifespan(fn=None)`
+
+Decorator and context manager to manage the lifecycle of a dependencies.
+This is equivalent of:
+
+```python
+import picodi
+
+
+picodi.init_dependencies()
+# your code
+picodi.shutdown_dependencies()
+```
+
+Can be used as a decorator:
+
+```python
+from picodi.helpers import lifespan
+
+
+@lifespan
+def main():
+    # Depedencies will be initialized before this function call
+    # and closed after this function call
+    ...
+
+# or for async functions
+@lifespan
+async def async_main():
+    ...
+```
+
+Can be used as a context manager:
+
+```python
+from picodi.helpers import lifespan
+
+
+with lifespan():  # or `async with lifespan():` for async functions
+    ...
+```
+
+`lifespan` can automatically detect if the function is async or not.
+But if you want to force sync or async mode,
+you can use `lifespan.sync` or `lifespan.async_`:
+
+```python
+from picodi.helpers import lifespan
+
+
+with lifespan.sync():
+    ...
+
+
+@lifespan.async_()
+async def main():
+    ...
+```
+
 ## License
 
 [MIT](https://github.com/yakimka/picodi/blob/main/LICENSE)

picodi/helpers.py

@@ -1,4 +1,14 @@
-from typing import Any
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, overload
+
+import picodi
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable, Generator
+    from types import TracebackType
 
 sentinel = object()
 
@@ -57,3 +67,67 @@ def _get_item(obj: Any, key: str) -> Any:
         return obj[key]
     except (KeyError, TypeError):
         return sentinel
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+class _Lifespan:
+    @overload
+    def __call__(self, fn: Callable[P, T]) -> Callable[P, T]:
+        pass
+
+    @overload
+    def __call__(self, fn: None = None) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        pass
+
+    def __call__(
+        self, fn: Callable[P, T] | None = None
+    ) -> Callable[P, T] | Callable[[Callable[P, T]], Callable[P, T]]:
+        if fn is None:
+            return self
+        if asyncio.iscoroutinefunction(fn):
+            return self.async_()(fn)  # type: ignore[return-value]
+        return self.sync()(fn)
+
+    def __enter__(self) -> None:
+        picodi.init_dependencies()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        picodi.shutdown_dependencies()
+
+    async def __aenter__(self) -> None:
+        await picodi.init_dependencies()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await picodi.shutdown_dependencies()
+
+    @contextlib.contextmanager
+    def sync(self) -> Generator[None, None, None]:
+        picodi.init_dependencies()
+        try:
+            yield
+        finally:
+            picodi.shutdown_dependencies()
+
+    @contextlib.asynccontextmanager
+    async def async_(self) -> AsyncGenerator[None, None]:
+        await picodi.init_dependencies()
+        try:
+            yield
+        finally:
+            await picodi.shutdown_dependencies()  # noqa: ASYNC102
+
+
+lifespan = _Lifespan()

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "picodi"
 description = "Simple Dependency Injection for Python"
-version = "0.13.0"
+version = "0.14.0"
 license = "MIT"
 authors = [
   "yakimka"
@@ -39,6 +39,7 @@ pytest-deadfixtures = "^2.2.1"
 pytest-randomly = "^3.12"
 fastapi-slim = "^0.111.0"
 httpx = "^0.27.0"
+pytest-markdown-docs = "^0.5.1"
 
 [tool.isort]
 # isort configuration:

tests/test_lifespan.py

@@ -0,0 +1,84 @@
+import pytest
+
+from picodi import SingletonScope, dependency
+from picodi.helpers import lifespan
+
+
+@pytest.fixture()
+def resource():
+    state = {"inited": False, "closed": False}
+
+    @dependency(scope_class=SingletonScope)
+    def my_resource():
+        state["inited"] = True
+        yield state
+        state["closed"] = True
+
+    return state, my_resource
+
+
+@pytest.fixture()
+def async_resource():
+    state = {"inited": False, "closed": False}
+
+    @dependency(scope_class=SingletonScope)
+    async def my_resource():
+        state["inited"] = True
+        yield state
+        state["closed"] = True
+
+    return state, my_resource
+
+
+@pytest.mark.parametrize("decorator", [lifespan, lifespan.sync()])
+def test_can_init_and_shutdown_sync(resource, decorator):
+    state, _ = resource
+
+    @decorator
+    def service():
+        assert state["inited"] is True
+        assert state["closed"] is False
+
+    service()
+
+    assert state["inited"] is True
+    assert state["closed"] is True
+
+
+@pytest.mark.parametrize("decorator", [lifespan, lifespan.async_()])
+async def test_can_init_and_shutdown_async(async_resource, decorator):
+    state, _ = async_resource
+
+    @decorator
+    async def service():
+        assert state["inited"] is True
+        assert state["closed"] is False
+
+    await service()
+
+    assert state["inited"] is True
+    assert state["closed"] is True
+
+
+@pytest.mark.parametrize("manager", [lifespan, lifespan.sync])
+def test_can_init_and_shutdown_sync_as_context_manager(resource, manager):
+    state, _ = resource
+
+    with manager():
+        assert state["inited"] is True
+        assert state["closed"] is False
+
+    assert state["inited"] is True
+    assert state["closed"] is True
+
+
+@pytest.mark.parametrize("manager", [lifespan, lifespan.async_])
+async def test_can_init_and_shutdown_async_as_context_manager(async_resource, manager):
+    state, _ = async_resource
+
+    async with manager():
+        assert state["inited"] is True
+        assert state["closed"] is False
+
+    assert state["inited"] is True
+    assert state["closed"] is True