Some tree_manager refactor and request scope contextvar fixes

Author: eadwinCode

ellar/di/__init__.py

@@ -7,10 +7,15 @@
 from .constants import (
     INJECTABLE_ATTRIBUTE,
     MODULE_REF_TYPES,
-    SCOPED_CONTEXT_VAR,
     AnnotationToValue,
+    request_context_var,
+)
+from .injector import (
+    Container,
+    EllarInjector,
+    ModuleTreeManager,
+    register_request_scope_context,
 )
-from .injector import Container, EllarInjector, ModuleTreeManager
 from .scopes import (
     RequestORTransientScope,
     RequestScope,
@@ -45,12 +50,13 @@
     "has_binding",
     "get_scope",
     "RequestScopeContext",
-    "SCOPED_CONTEXT_VAR",
+    "request_context_var",
     "INJECTABLE_ATTRIBUTE",
     "AnnotationToValue",
     "MODULE_REF_TYPES",
     "InjectByTag",
     "ModuleTreeManager",
+    "register_request_scope_context",
 ]
 
 

ellar/di/constants.py

@@ -6,10 +6,10 @@
 INJECTABLE_ATTRIBUTE = "__DI_SCOPE__"
 
 
-SCOPED_CONTEXT_VAR: contextvars.ContextVar[Optional[RequestScopeContext]] = (
-    contextvars.ContextVar("SCOPED-CONTEXT-VAR")
+request_context_var: contextvars.ContextVar[Optional[RequestScopeContext]] = (
+    contextvars.ContextVar("ellar.di.request_context_var")
 )
-SCOPED_CONTEXT_VAR.set(None)
+request_context_var.set(None)
 INJECTABLE_WATERMARK = "INJECTABLE_WATERMARK"
 
 
@@ -29,7 +29,7 @@ def __new__(mcls, name, bases, namespace):
         annotations.update(namespace.get("__annotations__", {}))
         cls.keys = []
         for k, v in annotations.items():
-            if type(v) == type(str):
+            if type(v) is type(str):
                 value = str(k).lower()
                 setattr(cls, k, value)
                 cls.keys.append(value)

ellar/di/context.py

@@ -1,5 +1,10 @@
 from .container import Container
-from .ellar_injector import EllarInjector
+from .ellar_injector import EllarInjector, register_request_scope_context
 from .tree_manager import ModuleTreeManager
 
-__all__ = ["Container", "EllarInjector", "ModuleTreeManager"]
+__all__ = [
+    "Container",
+    "EllarInjector",
+    "ModuleTreeManager",
+    "register_request_scope_context",
+]

ellar/di/injector/__init__.py

@@ -167,9 +167,6 @@ def get_binding(self, interface: t.Type) -> t.Tuple[Binding, InjectorBinder]:
             return super().get_binding(interface)
         except (KeyError, UnsatisfiedRequirement) as uex:
             try:
-                # All services or providers must have MODULE_SCOPE_OWNER if added in @Module 'providers',
-                # 'controllers' and 'routers'.
-                # So if 'module_scope_owner' is None, it means `Interface` was not registered to any module
                 if self.injector.owner:
                     module_name = (
                         self.injector.owner.name if self.injector.owner else None
@@ -184,23 +181,6 @@ def get_binding(self, interface: t.Type) -> t.Tuple[Binding, InjectorBinder]:
                         # TODO: possible circular import
                         return module_owner.value.container._get_binding(interface)
 
-                # current_node = self.injector.tree_manager.find_module(
-                #     predicate=lambda n: n.parent == self.injector.module_name
-                # )
-
-                # root_module_name = self.injector.tree_manager.get_root_module().name
-                #
-                # # module_data = self.injector.tree_manager.search_module_tree(
-                # #     filter_item=lambda data: data.name == module_name,
-                # #     find_predicate=lambda data: interface in data.exports,
-                # # )
-                # module_data = self.injector.tree_manager.search_module_tree(
-                #     filter_item=lambda data: data.name == module_name or data.name == root_module_name,
-                #     find_predicate=lambda data: interface in data.exports,
-                # )
-                # if module_data:
-                #     return module_data.value.container._get_binding(interface)
-
             except (KeyError, UnsatisfiedRequirement, Exception) as ex:
                 logging.exception(ex)
             raise UnsatisfiedRequirement(None, interface) from uex

ellar/di/injector/container.py

@@ -1,15 +1,13 @@
 import sys
 import typing as t
-from contextlib import asynccontextmanager
 from functools import cached_property
 
-from ellar.di.constants import MODULE_REF_TYPES, SCOPED_CONTEXT_VAR, Tag
+from ellar.di.constants import MODULE_REF_TYPES, Tag, request_context_var
 from ellar.di.injector.tree_manager import ModuleTreeManager
 from ellar.di.logger import log
 from injector import Injector, Scope, ScopeDecorator
 from typing_extensions import Annotated
 
-from ..asgi_args import RequestScopeContext
 from ..providers import InstanceProvider, Provider
 from ..types import T
 from .container import Container
@@ -22,6 +20,19 @@
     )
 
 
+def register_request_scope_context(interface: t.Type[T], value: T) -> None:
+    # Sets RequestScope contexts so that they can be available when needed
+
+    scoped_context = request_context_var.get()
+    if scoped_context is None:
+        return
+
+    if isinstance(value, Provider):
+        scoped_context.context.update({interface: value})
+    else:
+        scoped_context.context.update({interface: InstanceProvider(value)})
+
+
 class _TagInfo(t.NamedTuple):
     supertype: t.Type
     tag: str
@@ -124,23 +135,3 @@ def get(
         )
         log.debug(f"{self._log_prefix} -> {result}")
         return t.cast(T, result)
-
-    def update_scoped_context(self, interface: t.Type[T], value: T) -> None:
-        # Sets RequestScope contexts so that they can be available when needed
-        # TODO: Rename to `register_request_scope_context`
-        scoped_context = SCOPED_CONTEXT_VAR.get()
-        if scoped_context is None:
-            return
-
-        if isinstance(value, Provider):
-            scoped_context.context.update({interface: value})
-        else:
-            scoped_context.context.update({interface: InstanceProvider(value)})
-
-    @asynccontextmanager
-    async def create_asgi_args(self) -> t.AsyncGenerator["EllarInjector", None]:
-        try:
-            SCOPED_CONTEXT_VAR.set(RequestScopeContext())
-            yield self
-        finally:
-            SCOPED_CONTEXT_VAR.set(None)

ellar/di/injector/ellar_injector.py

@@ -30,8 +30,8 @@ def exports(self) -> t.List[t.Type]:
         return self.value.exports
 
     @property
-    def providers(self) -> t.Dict[t.Type, t.Type]:
-        return self.value.providers
+    def providers(self) -> t.Dict[t.Type, t.Union[t.Type, "ProviderConfig"]]:
+        return self.value.providers  # type:ignore[return-value]
 
     @property
     def name(self) -> str:
@@ -45,7 +45,7 @@ def __str__(self) -> str:
 
 
 class ModuleTreeManager:
-    __slots__ = ("modules", "_core_module", "_root_module")
+    __slots__ = ("modules", "_core_module", "_app_module")
 
     # , root_module: t.Union["ModuleRefBase", "ModuleSetup"]
     def __init__(
@@ -57,15 +57,16 @@ def __init__(
         )  # Dictionary to store modules by their ID or value
 
         self._core_module = app_core_module.module if app_core_module else None
-        self._root_module: t.Optional[t.Type[t.Any]] = None
+        self._app_module: t.Optional[t.Type[t.Any]] = None
 
         if app_core_module:
             self.add_module(app_core_module.module, value=app_core_module)
 
     @property
     def root_module(self) -> t.Type:
-        assert self._root_module is not None, "RootModule is not ready"
-        return self._root_module
+        root_module = self._core_module or self._app_module
+        assert root_module is not None, "RootModule is not ready"
+        return root_module
 
     def add_provider(
         self,
@@ -105,19 +106,19 @@ def add_module(
 
             self.modules[parent_module].dependencies.append(module_type)
 
-            if parent_module == self._core_module and not self._root_module:
-                self._root_module = data.value.module
+            if parent_module == self._core_module and not self._app_module:
+                self._app_module = data.value.module
 
-        elif not parent_module and not self._root_module and not self._core_module:
-            self._root_module = module_type
+        elif not parent_module and not self._app_module and not self._core_module:
+            self._app_module = module_type
         elif (
-            self._root_module
+            self._app_module
             and parent_module
             and self._core_module
             and parent_module == self._core_module
         ):
             raise Exception(
-                f"EllarCoreModule can only have '{self._root_module}' as dependency"
+                f"EllarCoreModule can only have '{self._app_module}' as dependency"
             )
         return self
 
@@ -204,8 +205,9 @@ def find_module(
         if not found_any:
             yield None  # type:ignore[misc]
 
-    def get_root_module(self) -> t.Union["ModuleRefBase", "ModuleSetup"]:
-        data = self.get_module(self.root_module)
+    def get_app_module(self) -> t.Union["ModuleRefBase", "ModuleSetup"]:
+        assert self._app_module is not None, "AppModule is not ready"
+        data = self.get_module(self._app_module)
         assert data
         return data.value
 

ellar/di/injector/tree_manager.py

@@ -1,6 +1,7 @@
+import logging
 import typing as t
 
-from ellar.di import SCOPED_CONTEXT_VAR, RequestScopeContext
+from ellar.di import RequestScopeContext, request_context_var
 from injector import (
     NoScope as TransientScope,
 )
@@ -23,7 +24,11 @@
 
 class RequestScope(InjectorScope):
     def get_context(self) -> t.Optional[RequestScopeContext]:
-        return SCOPED_CONTEXT_VAR.get()
+        try:
+            return request_context_var.get()
+        except Exception as ex:
+            logging.exception(ex)
+            return None
 
     def get(self, key: t.Type[T], provider: Provider[T]) -> Provider[T]:
         scoped_context = self.get_context()