Merge pull request #650 from reagento/fix/decorate_override

Check override works after decoration

Author: Tishka17

src/dishka/code_tools/factory_compiler.py

@@ -170,52 +170,53 @@ def _context_factory_body(
             ),
         )
 
-
 def _selector_factory_body(
     builder: FactoryBuilder, source_call: str, factory: Factory,
 ) -> None:
+    error_call = builder.call(
+        builder.global_(NoActiveFactoryError),
+        builder.global_(factory.provides),
+        builder.global_(factory.when_dependencies, "when_dependencies"),
+    )
+    builder.raise_(error_call)
+
+ASYNC_TYPES = (FactoryType.ASYNC_FACTORY, FactoryType.ASYNC_GENERATOR)
+BODY_GENERATORS = {
+    FactoryType.FACTORY: _sync_factory_body,
+    FactoryType.ASYNC_FACTORY: _async_factory_body,
+    FactoryType.GENERATOR: _generator_body,
+    FactoryType.ASYNC_GENERATOR: _async_generator_body,
+    FactoryType.CONTEXT: _context_factory_body,
+    FactoryType.VALUE: _value_factory_body,
+    FactoryType.ALIAS: _alias_factory_body,
+    FactoryType.SELECTOR: _selector_factory_body,
+}
+
+
+def _select_when_dependency(
+    builder: FactoryBuilder,
+    factory: Factory,
+) -> bool:
+    """return True if there is assignment in any case"""
     first = True
     for variant in factory.when_dependencies:
         condition = builder.when(variant.when_override, factory.when_component)
         solved_value = builder.getter(variant.provides)
         if first and not condition:
             builder.assign_solved(solved_value)
+            return True
         elif first:
             with builder.if_(condition):
                 builder.assign_solved(solved_value)
-            first = False
         elif not condition:
             with builder.else_():
                 builder.assign_solved(solved_value)
-            first = True
+            return True
         else:
             with builder.elif_(condition):
                 builder.assign_solved(solved_value)
-    # if-chain not closed with else or not generated at all
-    if not first or not factory.when_dependencies:
-        error_call = builder.call(
-            builder.global_(NoActiveFactoryError),
-            builder.global_(factory.provides),
-            builder.global_(factory.when_dependencies, "when_dependencies"),
-        )
-        if first:  # no options at all
-            builder.raise_(error_call)
-        else:
-            with builder.else_():
-                builder.raise_(error_call)
-
-
-ASYNC_TYPES = (FactoryType.ASYNC_FACTORY, FactoryType.ASYNC_GENERATOR)
-BODY_GENERATORS = {
-    FactoryType.FACTORY: _sync_factory_body,
-    FactoryType.ASYNC_FACTORY: _async_factory_body,
-    FactoryType.GENERATOR: _generator_body,
-    FactoryType.ASYNC_GENERATOR: _async_generator_body,
-    FactoryType.CONTEXT: _context_factory_body,
-    FactoryType.VALUE: _value_factory_body,
-    FactoryType.ALIAS: _alias_factory_body,
-    FactoryType.SELECTOR: _selector_factory_body,
-}
+        first = False
+    return False
 
 
 def compile_factory(*, factory: Factory, is_async: bool) -> CompiledFactory:
@@ -228,16 +229,22 @@ def compile_factory(*, factory: Factory, is_async: bool) -> CompiledFactory:
     builder.register_provides(factory.provides)
 
     with builder.make_getter():
-        source_call = builder.call(
-            builder.global_(factory.source),
-            *(builder.getter(dep) for dep in factory.dependencies),
-            **{
-                name: builder.getter(dep)
-                for name, dep in factory.kw_dependencies.items()
-            },
-        )
-        body_generator = BODY_GENERATORS[factory.type]
-        body_generator(builder, source_call, factory)
+        if not _select_when_dependency(builder, factory):
+            source_call = builder.call(
+                builder.global_(factory.source),
+                *(builder.getter(dep) for dep in factory.dependencies),
+                **{
+                    name: builder.getter(dep)
+                    for name, dep in factory.kw_dependencies.items()
+                },
+            )
+            body_generator = BODY_GENERATORS[factory.type]
+            if factory.when_dependencies:  # conditions generated
+                with builder.else_():
+                    body_generator(builder, source_call, factory)
+            else:  # no options at all
+                body_generator(builder, source_call, factory)
+
         if factory.cache:
             builder.cache()
         builder.return_("solved")

src/dishka/dependency_source/decorator.py

@@ -42,8 +42,6 @@ def as_factory(
             new_dependency: DependencyKey,
             cache: bool,
             component: Component,
-            when_override: BaseMarker | None,
-            when_active: BaseMarker | None,
     ) -> Factory:
         typevar_replacement = get_typevar_replacement(
             self.provides.type_hint,
@@ -71,8 +69,8 @@ def as_factory(
             },
             type_=self.factory.type,
             cache=cache,
-            when_override=combine_when(self.when, when_override),
-            when_active=combine_when(self.when, when_active),
+            when_override=self.when,
+            when_active=self.when,
             when_component=self.factory.when_component or component,
             when_dependencies=[],
         )

src/dishka/dependency_source/factory.py

@@ -4,7 +4,8 @@
     Mapping,
     Sequence,
 )
-from typing import Any
+from enum import Enum
+from typing import Any, TypeAlias, TypeVar
 
 from dishka.entities.component import Component
 from dishka.entities.factory_type import FactoryData, FactoryType
@@ -13,6 +14,19 @@
 from dishka.entities.scope import BaseScope
 
 
+class Special(Enum):
+    OMITTED = "omitted"
+
+
+T = TypeVar("T")
+MayBe: TypeAlias = Special | T
+
+def coalesce(a: MayBe[T], b: T) -> T:
+    if a is Special.OMITTED:
+        return b
+    return a
+
+
 class Factory(FactoryData):
     __slots__ = (
         "cache",
@@ -140,19 +154,22 @@ def with_scope(self, scope: BaseScope) -> Factory:
 
     def replace(
         self,
-        provides: DependencyKey | None = None,
+        provides: MayBe[DependencyKey] = Special.OMITTED,
+        when_active: MayBe[BaseMarker|None] = Special.OMITTED,
+        when_override: MayBe[BaseMarker|None] = Special.OMITTED,
+        when_component: MayBe[Component] = Special.OMITTED,
     ) -> Factory:
         return Factory(
             dependencies=list(self.dependencies),
             kw_dependencies=dict(self.kw_dependencies),
             source=self.source,
-            provides=provides or self.provides,
+            provides=coalesce(provides, self.provides),
             scope=self.scope,
             is_to_bind=self.is_to_bind,
             cache=self.cache,
             type_=self.type,
-            when_override=self.when_override,
-            when_active=self.when_active,
-            when_component=self.when_component,
+            when_override=coalesce(when_override, self.when_override),
+            when_active=coalesce(when_active, self.when_active),
+            when_component=coalesce(when_component, self.when_component),
             when_dependencies=self.when_dependencies,
         )

src/dishka/registry_builder.py

@@ -229,7 +229,6 @@ def _ensure_override_flags(
     ) -> None:
         if self.skip_validation:
             return
-
         if (
             not prev_factory and
             self.validation_settings.nothing_overridden and
@@ -253,8 +252,10 @@ def _unite_factories_group(
         group: list[Factory],
     ) -> dict[DependencyKey, list[Factory]]:
         if len(group) == 1:
-            self._ensure_override_flags(group[0], None)
-            return {}
+            factory = group[0]
+            self._ensure_override_flags(factory, None)
+            if factory.when_override in (None, BoolMarker(True)):
+                return {}
 
         when_dependencies: list[Factory] = []
         moved_factories: dict[DependencyKey, list[Factory]] = {}
@@ -279,10 +280,12 @@ def _unite_factories_group(
             new_factory = factory.replace(provides=new_provides)
             moved_factories[new_provides] = [new_factory]
             when_dependencies.append(new_factory)
-        if len(moved_factories) == 1:
-            self.processed_factories[provides] = [
-                cast(Factory, prev_factory),  # at least one factory found
-            ]
+        if (
+            len(moved_factories) == 1 and
+            prev_factory and  # at least one factory found
+            prev_factory.when_override in (None, BoolMarker(True))
+        ):
+            self.processed_factories[provides] = [prev_factory]
             return {}
 
         scope = max(
@@ -435,8 +438,6 @@ def _process_generic_decorator(
                     new_dependency=provides,
                     cache=False,
                     component=provider.component,
-                    when_override=None,
-                    when_active=None,
                 )],
             )
 
@@ -456,8 +457,6 @@ def _process_normal_decorator(
                     new_dependency=provides,
                     cache=False,
                     component=provider.component,
-                    when_active=None,
-                    when_override=None,
                 )],
             )
         self._decorate_factory(
@@ -497,8 +496,6 @@ def _decorate_factory(
                 f"no factory for {provides}",
             )
 
-        if decorator.when not in (None, BoolMarker(True)):
-            group_replacement.extend(old_group)
         for old_factory in old_group:
 
             depth = self.decorator_depth[provides]
@@ -515,18 +512,32 @@ def _decorate_factory(
             ):
                 return
 
-            new_factory = old_factory.replace(provides=decorated_provides)
+            new_factory = old_factory.replace(
+                provides=decorated_provides,
+                when_active=None,
+                when_override=None,
+            )
             decorated_groups[decorated_provides] = [new_factory]
             decorated_factory = decorator.as_factory(
                 scope=cast(BaseScope, old_factory.scope),
                 new_dependency=decorated_provides,
                 cache=old_factory.cache,
                 component=provides.component,
-                when_override=old_factory.when_override,
+            ).replace(
+                provides=provides,
                 when_active=old_factory.when_active,
-            ).replace(provides=provides)
+                when_override=old_factory.when_override,
+                when_component=cast(Component, old_factory.when_component),
+            )
+            if decorator.when is not None:
+                conditional_factory = new_factory.replace(
+                    when_override=~decorator.when,
+                    when_active=~decorator.when,
+                    when_component=provides.component,
+                )
+                decorated_factory.when_dependencies=[conditional_factory]
+                self._register_when(conditional_factory)
             group_replacement.append(decorated_factory)
-            self._register_when(decorated_factory)
 
         self.processed_factories[provides] = group_replacement
         self.processed_factories.update(decorated_groups)

tests/unit/container/test_decorator.py

@@ -4,6 +4,7 @@
 
 from dishka import (
     DEFAULT_COMPONENT,
+    STRICT_VALIDATION,
     DependencyKey,
     Has,
     Provider,
@@ -13,7 +14,11 @@
     make_container,
     provide,
 )
-from dishka.exceptions import CycleDependenciesError, NoFactoryError
+from dishka.exceptions import (
+    CycleDependenciesError,
+    ImplicitOverrideDetectedError,
+    NoFactoryError,
+)
 
 
 class A:
@@ -398,3 +403,45 @@ def dec(self, a: A) -> A:
 
     with pytest.raises(NoFactoryError):
         make_container(MyProvider())
+
+
+def test_decorate_override():
+    class MyProvider(Provider):
+        @provide
+        def make_str(self) -> str:
+            return "a"
+
+        @provide(override=True)
+        def make_str2(self) -> str:
+            return "b"
+
+        @decorate
+        def decorate_str(self, old_value: str) -> str:
+            return old_value + "d"
+
+    c = make_container(
+        MyProvider(scope=Scope.APP),
+        validation_settings=STRICT_VALIDATION,
+    )
+    assert c.get(str) == "bd"
+
+
+def test_decorate_override_implicit():
+    class MyProvider(Provider):
+        @provide
+        def make_str(self) -> str:
+            return "a"
+
+        @provide
+        def make_str2(self) -> str:
+            return "b"
+
+        @decorate
+        def decorate_str(self, old_value: str) -> str:
+            return old_value + "d"
+
+    with pytest.raises(ImplicitOverrideDetectedError):
+        make_container(
+            MyProvider(scope=Scope.APP),
+            validation_settings=STRICT_VALIDATION,
+        )

tests/unit/container/test_exceptions.py

@@ -335,11 +335,20 @@ def test_no_active_factory_smoke(path_len: int, variants_count: int)  -> None:
 
 
 @pytest.mark.parametrize("is_async", [True, False])
+@pytest.mark.parametrize(("has_variants", "variants_count"), [
+    ([float], 1),
+    ([float, complex], 2),
+])
 @pytest.mark.asyncio
-async def test_no_active_factory(*, is_async: bool) -> None:
+async def test_no_active_factory(
+    *,
+    is_async: bool,
+    has_variants: list[object],
+    variants_count: int,
+) -> None:
     provider = Provider(scope=Scope.APP)
-    provider.provide(int, when=Has(float))
-    provider.provide(int, when=Has(complex))
+    for variant in has_variants:
+        provider.provide(int, when=Has(variant))
 
     @provider.provide
     def get_str(value: int) -> str:
@@ -356,5 +365,5 @@ def get_str(value: int) -> str:
 
 
     assert str(e.value)
-    assert len(e.value.variants) == 2
+    assert len(e.value.variants) == variants_count
     assert len(e.value.path) == 2