Add Async Support

minject/inject.py

@@ -2,22 +2,36 @@
 
 import itertools
 import os
-from typing import Any, Callable, Dict, Optional, Sequence, Type, TypeVar, Union, cast, overload
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 from typing_extensions import TypeGuard
 
-from .metadata import RegistryMetadata, _gen_meta, _get_meta
+from minject.types import _AsyncContext
+
+from .metadata import _INJECT_METADATA_ATTR, RegistryMetadata, _gen_meta, _get_meta
 from .model import (
     Deferred,
     DeferredAny,
     RegistryKey,  # pylint: disable=unused-import
     Resolver,
     resolve_value,
 )
-from .types import _MinimalMappingProtocol
+from .types import _AsyncContext, _MinimalMappingProtocol
 
 T = TypeVar("T")
 T_co = TypeVar("T_co", covariant=True)
+T_async_context = TypeVar("T_async_context", bound=_AsyncContext)
 R = TypeVar("R")
 
 
@@ -70,6 +84,19 @@ def wrap(cls: Type[T]) -> Type[T]:
     return wrap
 
 
+def async_context(cls: Type[T_async_context]) -> Type[T_async_context]:
+    """
+    Decorator to declare that a class is as an async context manager
+    that can be initialized by the registry through aget(). This
+    is to distinguish the class from an async context manager that
+    should not be initialized by the registry (an example of
+    this being asyncio.Lock()).
+    """
+    meta = _gen_meta(cls)
+    meta.update_async_context(True)
+    return cls
+
+
 def define(
     base_class: Type[T],
     _close: Optional[Callable[[T], None]] = None,
@@ -91,6 +118,24 @@ def _is_type(key: "RegistryKey[T]") -> TypeGuard[Type[T]]:
     return isinstance(key, type)
 
 
+def _is_key_async(key: RegistryKey) -> bool:
+    """
+    Check whether a registry key is an async context manager
+    marked for initialization within the registry (@async_context).
+    """
+    if isinstance(key, str):
+        return False
+    elif isinstance(key, RegistryMetadata):
+        return key.is_async_context()
+    elif isinstance(key, type):
+        try:  # type: ignore
+            inject_metadata = object.__getattribute__(key, _INJECT_METADATA_ATTR)
+            return inject_metadata.is_async_context()  # type: ignore
+        except AttributeError:
+            return False
+    assert False, f"Unexpected key type: {key}"
+
+
 class _RegistryReference(Deferred[T_co]):
     """Reference to an object in the registry to be loaded later.
     (you should not instantiate this class directly, instead use the
@@ -103,6 +148,13 @@ def __init__(self, key: "RegistryKey[T_co]") -> None:
     def resolve(self, registry_impl: Resolver) -> T_co:
         return registry_impl.resolve(self._key)
 
+    async def aresolve(self, registry_impl: Resolver) -> T_co:
+        if _is_key_async(self._key):
+            result = await registry_impl.aresolve(self._key)
+            await registry_impl.push_async_context(result)
+            return result
+        return registry_impl.resolve(self._key)
+
     @property
     def type_of_object_referenced_in_key(self) -> "Type[T_co]":
         if type(self.key) == RegistryMetadata:
@@ -188,6 +240,9 @@ def resolve(self, registry_impl: Resolver) -> T_co:
             kwargs[key] = resolve_value(registry_impl, arg)
         return self.func()(*args, **kwargs)
 
+    async def aresolve(self, registry_impl: Resolver) -> T_co:
+        raise NotImplementedError("Have not implemented async registry function")
+
     def func(self) -> Callable[..., T_co]:
         return self._func
 
@@ -273,6 +328,9 @@ def resolve(self, registry_impl: Resolver) -> T_co:
         else:
             return cast(T_co, self._default)
 
+    async def aresolve(self, registry_impl: Resolver) -> T_co:
+        return self.resolve(registry_impl)
+
     @property
     def key(self) -> Optional[str]:
         return self._key
@@ -316,6 +374,9 @@ def resolve(self, registry_impl: Resolver) -> T_co:
                 return self._default
         return cast(T_co, sub)
 
+    async def aresolve(self, registry_impl: Resolver) -> T_co:
+        return self.resolve(registry_impl)
+
 
 def nested_config(
     keys: Union[Sequence[str], str], default: Union[T, _RaiseKeyError] = RAISE_KEY_ERROR
@@ -375,5 +436,8 @@ class _RegistrySelf(Deferred[Resolver]):
     def resolve(self, registry_impl: Resolver) -> Resolver:
         return registry_impl
 
+    async def aresolve(self, registry_impl: Resolver) -> Resolver:
+        return self.resolve(registry_impl)
+
 
 self_tag = _RegistrySelf()

minject/metadata.py

@@ -112,6 +112,7 @@ def __init__(
 
         self._close = close
         self._interfaces = [cls for cls in inspect.getmro(cls) if cls is not object]
+        self._is_async_context = False
 
     @property
     def interfaces(self) -> Sequence[Type]:
@@ -142,6 +143,18 @@ def update_bindings(self, **bindings: DeferredAny) -> None:
         # TODO: 'lock' the bindings once added to the registry to make above note unnecessary
         self._bindings.update(bindings)
 
+    def update_async_context(self, is_async_context: bool) -> None:
+        """
+        Set the async context flag for this metadata.
+        """
+        self._is_async_context = is_async_context
+
+    def is_async_context(self) -> bool:
+        """
+        Returns the value of the async context flag for this metadata.
+        """
+        return self._is_async_context
+
     def _new_object(self) -> T_co:
         return self._cls.__new__(self._cls)
 
@@ -159,6 +172,17 @@ def _init_object(self, obj: T_co, registry_impl: "Registry") -> None:  # type: i
 
         self._cls.__init__(obj, **init_kwargs)
 
+    async def _ainit_object(self, obj: T_co, registry_impl: "Registry") -> None:  # type: ignore[misc]
+        """
+        asynchronous version of _init_object. Calls _aresolve instead
+        of _resolve.
+        """
+        init_kwargs = {}
+        for name_, value in self._bindings.items():
+            # the specific deferred value checks if they are async or
+            init_kwargs[name_] = await registry_impl._aresolve(value)
+        self._cls.__init__(obj, **init_kwargs)
+
     def _close_object(self, obj: T_co) -> None:  # type: ignore[misc]
         if self._close:
             self._close(obj)

minject/model.py

@@ -34,6 +34,23 @@ class Resolver(abc.ABC):
     def resolve(self, key: "RegistryKey[T]") -> T:
         ...
 
+    @abc.abstractmethod
+    async def aresolve(self, key: "RegistryKey[T]") -> T:
+        """
+        Resolve a key into an instance of that key. The key must be marked
+        with the @async_context decorator.
+        """
+        ...
+
+    @abc.abstractmethod
+    async def push_async_context(self, key: Any) -> Any:
+        """
+        Push an async context onto the context stack maintained by the Resolver.
+        This is necessary to enter/close the context of an object
+        marked with @async_context.
+        """
+        ...
+
     @property
     @abc.abstractmethod
     def config(self) -> RegistryConfigWrapper:
@@ -52,6 +69,15 @@ class Deferred(abc.ABC, Generic[T_co]):
     def resolve(self, registry_impl: Resolver) -> T_co:
         ...
 
+    @abc.abstractmethod
+    async def aresolve(self, registry_impl: Resolver) -> T_co:
+        """
+        Resolve a deferred object into an instance of the object. The object,
+        may or may not be asynchronous. If the object is asynchronous (marked with @async_context),
+        resolve the object asynchronously. Otherwise, resolve synchronously.
+        """
+        ...
+
 
 Resolvable = Union[Deferred[T_co], T_co]
 # Union of Deferred and Any is just Any, but want to call out that a Deffered is quite common
@@ -69,3 +95,13 @@ def resolve_value(registry_impl: Resolver, value: Resolvable[T]) -> T:
         return value.resolve(registry_impl)
     else:
         return value
+
+
+async def aresolve_value(registry_impl: Resolver, value: Resolvable[T]) -> T:
+    """
+    Async version of resolve_value, which calls aresolve on Deferred instances.
+    """
+    if isinstance(value, Deferred):
+        return await value.aresolve(registry_impl)
+    else:
+        return value