wip

Author: remimd

injection/__init__.py

@@ -1,19 +1,18 @@
 from ._core.descriptors import LazyInstance
 from ._core.injectables import Injectable
 from ._core.module import Mode, Module, Priority, mod
-from ._core.scope import AsyncScope, SyncScope
+from ._core.scope import async_scope, sync_scope
 
 __all__ = (
-    "AsyncScope",
     "Injectable",
     "LazyInstance",
     "Mode",
     "Module",
     "Priority",
-    "SyncScope",
     "afind_instance",
     "aget_instance",
     "aget_lazy_instance",
+    "async_scope",
     "constant",
     "find_instance",
     "get_instance",
@@ -25,6 +24,7 @@
     "set_constant",
     "should_be_injectable",
     "singleton",
+    "sync_scope",
 )
 
 afind_instance = mod().afind_instance

injection/__init__.pyi

@@ -1,11 +1,11 @@
 from abc import abstractmethod
 from collections.abc import Awaitable, Callable
-from contextlib import AsyncContextDecorator, ContextDecorator
+from contextlib import ContextDecorator
 from enum import Enum
 from logging import Logger
-from types import TracebackType
 from typing import (
     Any,
+    AsyncContextManager,
     ContextManager,
     Final,
     Protocol,
@@ -37,6 +37,8 @@ set_constant = __MODULE.set_constant
 should_be_injectable = __MODULE.should_be_injectable
 singleton = __MODULE.singleton
 
+def async_scope(name: str, *, shared: bool = ...) -> AsyncContextManager[None]: ...
+def sync_scope(name: str, *, shared: bool = ...) -> ContextManager[None]: ...
 def mod(name: str = ..., /) -> Module:
     """
     Short syntax for `Module.from_name`.
@@ -318,23 +320,3 @@ class Mode(Enum):
 class Priority(Enum):
     LOW = ...
     HIGH = ...
-
-class AsyncScope(AsyncContextDecorator):
-    def __init__(self, name: str) -> None: ...
-    async def __aenter__(self) -> Self: ...
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> Any: ...
-
-class SyncScope(ContextDecorator):
-    def __init__(self, name: str) -> None: ...
-    def __enter__(self) -> Self: ...
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> Any: ...

injection/_core/scope.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import Iterator, MutableMapping
+from collections.abc import AsyncIterator, Iterator, MutableMapping
 from contextlib import (
     AsyncContextDecorator,
     AsyncExitStack,
     ContextDecorator,
     ExitStack,
+    asynccontextmanager,
     contextmanager,
 )
 from contextvars import ContextVar
@@ -31,59 +32,111 @@
 )
 
 
-@dataclass(repr=False, frozen=True, slots=True)
-class _ActiveScope:
+@dataclass(repr=False, slots=True)
+class _ScopeState:
     # Shouldn't be instantiated outside `__SCOPES`.
 
-    context_var: ContextVar[Scope] = field(
+    __context_var: ContextVar[Scope] = field(
         default_factory=lambda: ContextVar(f"scope@{new_short_key()}"),
         init=False,
     )
-    references: set[Scope] = field(
+    __references: set[Scope] = field(
         default_factory=set,
         init=False,
     )
+    __shared_value: Scope | None = field(
+        default=None,
+        init=False,
+    )
 
-    def to_tuple(self) -> tuple[ContextVar[Scope], set[Scope]]:
-        return self.context_var, self.references
+    @contextmanager
+    def bind_contextual_scope(self, scope: Scope) -> Iterator[None]:
+        self.__references.add(scope)
+        token = self.__context_var.set(scope)
 
+        try:
+            yield
+        finally:
+            self.__context_var.reset(token)
+            self.__references.remove(scope)
 
-__SCOPES: Final[defaultdict[str, _ActiveScope]] = defaultdict(_ActiveScope)
+    @contextmanager
+    def bind_shared_scope(self, scope: Scope) -> Iterator[None]:
+        if self.__references:
+            raise ScopeError(
+                "A shared scope can't be defined when one or more contextual scopes "
+                "are defined on the same name."
+            )
 
+        self.__shared_value = scope
 
-@contextmanager
-def bind_scope(name: str, value: Scope) -> Iterator[None]:
-    context_var, references = __SCOPES[name].to_tuple()
+        try:
+            yield
+        finally:
+            self.__shared_value = None
 
-    if context_var.get(None):
-        raise ScopeAlreadyDefinedError(
-            f"Scope `{name}` is already defined in the current context."
-        )
+    def get_scope(self) -> Scope | None:
+        return self.__context_var.get(self.__shared_value)
 
-    references.add(value)
-    token = context_var.set(value)
+    def get_active_scopes(self) -> tuple[Scope, ...]:
+        references = self.__references
 
-    try:
+        if shared_value := self.__shared_value:
+            return shared_value, *references
+
+        return tuple(references)
+
+
+__SCOPES: Final[defaultdict[str, _ScopeState]] = defaultdict(_ScopeState)
+
+
+@asynccontextmanager
+async def async_scope(name: str, *, shared: bool = False) -> AsyncIterator[None]:
+    async with AsyncScope() as scope:
+        scope.enter(_bind_scope(name, scope, shared))
+        yield
+
+
+@contextmanager
+def sync_scope(name: str, *, shared: bool = False) -> Iterator[None]:
+    with SyncScope() as scope:
+        scope.enter(_bind_scope(name, scope, shared))
         yield
-    finally:
-        context_var.reset(token)
-        references.discard(value)
-        value.cache.clear()
 
 
 def get_active_scopes(name: str) -> tuple[Scope, ...]:
-    return tuple(__SCOPES[name].references)
+    return __SCOPES[name].get_active_scopes()
 
 
 def get_scope(name: str) -> Scope:
-    context_var = __SCOPES[name].context_var
+    scope = __SCOPES[name].get_scope()
 
-    try:
-        return context_var.get()
-    except LookupError as exc:
+    if scope is None:
         raise ScopeUndefinedError(
             f"Scope `{name}` isn't defined in the current context."
-        ) from exc
+        )
+
+    return scope
+
+
+@contextmanager
+def _bind_scope(name: str, value: Scope, shared: bool) -> Iterator[None]:
+    state = __SCOPES[name]
+
+    if state.get_scope():
+        raise ScopeAlreadyDefinedError(
+            f"Scope `{name}` is already defined in the current context."
+        )
+
+    strategy = (
+        state.bind_shared_scope(value) if shared else state.bind_contextual_scope(value)
+    )
+
+    try:
+        with strategy:
+            yield
+    finally:
+        value.cache.clear()
 
 
 @runtime_checkable
@@ -102,22 +155,23 @@ def enter[T](self, context_manager: ContextManager[T]) -> T:
 
 
 @dataclass(repr=False, frozen=True, slots=True)
-class AsyncScope(AsyncContextDecorator, Scope):
-    name: str
+class BaseScope[T](Scope, ABC):
+    delegate: T
     cache: MutableMapping[Any, Any] = field(
         default_factory=dict,
         init=False,
         hash=False,
     )
-    __exit_stack: AsyncExitStack = field(
-        default_factory=AsyncExitStack,
-        init=False,
-    )
+
+
+class AsyncScope(AsyncContextDecorator, BaseScope[AsyncExitStack]):
+    __slots__ = ()
+
+    def __init__(self) -> None:
+        super().__init__(delegate=AsyncExitStack())
 
     async def __aenter__(self) -> Self:
-        await self.__exit_stack.__aenter__()
-        lifespan = bind_scope(self.name, self)
-        self.enter(lifespan)
+        await self.delegate.__aenter__()
         return self
 
     async def __aexit__(
@@ -126,32 +180,23 @@ async def __aexit__(
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> Any:
-        return await self.__exit_stack.__aexit__(exc_type, exc_value, traceback)
+        return await self.delegate.__aexit__(exc_type, exc_value, traceback)
 
     async def aenter[T](self, context_manager: AsyncContextManager[T]) -> T:
-        return await self.__exit_stack.enter_async_context(context_manager)
+        return await self.delegate.enter_async_context(context_manager)
 
     def enter[T](self, context_manager: ContextManager[T]) -> T:
-        return self.__exit_stack.enter_context(context_manager)
+        return self.delegate.enter_context(context_manager)
 
 
-@dataclass(repr=False, frozen=True, slots=True)
-class SyncScope(ContextDecorator, Scope):
-    name: str
-    cache: MutableMapping[Any, Any] = field(
-        default_factory=dict,
-        init=False,
-        hash=False,
-    )
-    __exit_stack: ExitStack = field(
-        default_factory=ExitStack,
-        init=False,
-    )
+class SyncScope(ContextDecorator, BaseScope[ExitStack]):
+    __slots__ = ()
+
+    def __init__(self) -> None:
+        super().__init__(delegate=ExitStack())
 
     def __enter__(self) -> Self:
-        self.__exit_stack.__enter__()
-        lifespan = bind_scope(self.name, self)
-        self.enter(lifespan)
+        self.delegate.__enter__()
         return self
 
     def __exit__(
@@ -160,10 +205,10 @@ def __exit__(
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> Any:
-        return self.__exit_stack.__exit__(exc_type, exc_value, traceback)
+        return self.delegate.__exit__(exc_type, exc_value, traceback)
 
     async def aenter[T](self, context_manager: AsyncContextManager[T]) -> T:
         raise ScopeError("SyncScope doesn't support asynchronous context manager.")
 
     def enter[T](self, context_manager: ContextManager[T]) -> T:
-        return self.__exit_stack.enter_context(context_manager)
+        return self.delegate.enter_context(context_manager)

tests/core/test_module.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from injection import Module, SyncScope
+from injection import Module, sync_scope
 from injection.exceptions import (
     ModuleError,
     ModuleLockError,
@@ -346,7 +346,7 @@ class Dependency: ...
 
         assert module.is_locked is False
 
-        with SyncScope("test"):
+        with sync_scope("test"):
             instance_1 = module.get_instance(Dependency)
             assert module.is_locked is True