Repository: reagento/dishka
Pull request #268 (2024-10-14 21:45:00)

Simplify plots


Patch:
diff --git a/examples/integrations/click/__init__.py b/examples/integrations/click_app/__init__.py
similarity index 100%
rename from examples/integrations/click/__init__.py
rename to examples/integrations/click_app/__init__.py
diff --git a/examples/integrations/click/async_command.py b/examples/integrations/click_app/async_command.py
similarity index 100%
rename from examples/integrations/click/async_command.py
rename to examples/integrations/click_app/async_command.py
diff --git a/examples/integrations/click/sync_command.py b/examples/integrations/click_app/sync_command.py
similarity index 100%
rename from examples/integrations/click/sync_command.py
rename to examples/integrations/click_app/sync_command.py
diff --git a/requirements/faststream-047.txt b/requirements/faststream-047.txt
index 3a6047b..998b329 100644
--- a/requirements/faststream-047.txt
+++ b/requirements/faststream-047.txt
@@ -2,3 +2,4 @@
 
 faststream[nats]==0.4.7
 typing-extensions==4.11.0
+anyio==4.6.2
diff --git a/requirements/faststream-050.txt b/requirements/faststream-050.txt
index 321f5b0..7a2b11e 100644
--- a/requirements/faststream-050.txt
+++ b/requirements/faststream-050.txt
@@ -2,3 +2,4 @@
 
 faststream[nats]==0.5.0rc2
 typing-extensions==4.11.0
+anyio==4.6.2
diff --git a/src/dishka/entities/scope.py b/src/dishka/entities/scope.py
index 994d067..7548fe7 100644
--- a/src/dishka/entities/scope.py
+++ b/src/dishka/entities/scope.py
@@ -34,4 +34,4 @@ class InvalidScopes(BaseScope):
     UNKNOWN_SCOPE = new_scope("<unknown scope>", skip=True)
 
     def __str__(self) -> Any:
-        return str(self.value.name)
+        return str(self.value)
diff --git a/src/dishka/plotter/d2.py b/src/dishka/plotter/d2.py
index 3e289cc..6ebbc7c 100644
--- a/src/dishka/plotter/d2.py
+++ b/src/dishka/plotter/d2.py
@@ -4,14 +4,16 @@ from dishka.plotter.model import Group, Node, NodeType, Renderer
 
 
 class D2Renderer(Renderer):
-    def __init__(self):
-        self.names: dict[str, str] = {}
+    def __init__(self) -> None:
+        self.nodes: dict[str, Node] = {}
         self.full_ids: dict[str, str] = {}
 
     def _escape(self, line: str) -> str:
         return re.sub(r"[^\w_.\-]", "_", line)
 
     def _render_node(self, node: Node) -> str:
+        if node.type is NodeType.ALIAS:
+            return ""
         name = self._node_type(node) + self._escape(node.name)
         if node.source_name:
             source = f'    "{self._escape(node.source_name)}()": ""\n'
@@ -22,26 +24,37 @@ class D2Renderer(Renderer):
         res += "    shape: class\n"
         res += source
         for dep in node.dependencies:
-            dep_name = self._escape(self.names[dep])
+            dep_name = self._escape(self.nodes[dep].name)
             res += f"    {dep_name}\n"
         res += "}\n"
         return res
 
     def _render_node_deps(self, node: Node) -> list[str]:
-        return [
-            f"{self.full_ids[dep]} --> {self.full_ids[node.id]}"
-            for dep in node.dependencies
-        ]
-
-    def _group_type(self, group: Group) -> str:
-        return ""
+        res: list[str] = []
+        if node.type is NodeType.ALIAS:
+            return res
+        node_id = self.full_ids[node.id]
+        for dep in node.dependencies:
+            dep_node = self.nodes[dep]
+            if dep_node.type is NodeType.ALIAS:
+                target = self.full_ids[dep_node.dependencies[0]]
+                res.append(
+                    f"{target} --> {node_id}: "
+                    f"🔗 {self._escape(dep_node.name)}",
+                )
+            else:
+                target = self.full_ids[dep]
+                res.append(f"{target} --> {node_id}")
+        return res
 
     def _node_type(self, node: Node) -> str:
         if node.is_protocol:
             prefix = "Ⓘ "
         else:
             prefix = ""
-        if node.type is NodeType.CONTEXT:
+        if node.type is NodeType.DECORATOR:
+            return "🎭 " + prefix
+        elif node.type is NodeType.CONTEXT:
             return "📥 " + prefix
         elif node.type is NodeType.ALIAS:
             return "🔗 " + prefix
@@ -50,11 +63,7 @@ class D2Renderer(Renderer):
     def _render_group(
             self, group: Group, name_prefix: str = "",
     ) -> str:
-        name = (
-            self._group_type(group)
-            + name_prefix
-            + self._escape(group.name or "")
-        )
+        name = name_prefix + self._escape(group.name or "")
         res = f'{group.id}: "{name}"' + "{\n"
         for node in group.nodes:
             res += self._render_node(node) + "\n"
@@ -80,7 +89,7 @@ class D2Renderer(Renderer):
                 id_prefix = group.id
             self.full_ids[group.id] = id_prefix
             for node in group.nodes:
-                self.names[node.id] = node.name
+                self.nodes[node.id] = node
                 self.full_ids[node.id] = id_prefix + "." + node.id
             self._fill_names(group.children, prefix=id_prefix)
 
diff --git a/src/dishka/plotter/mermaid.py b/src/dishka/plotter/mermaid.py
index 354af89..272b0f3 100644
--- a/src/dishka/plotter/mermaid.py
+++ b/src/dishka/plotter/mermaid.py
@@ -25,17 +25,19 @@ HTML_TEMPLATE = """\
 
 
 class MermaidRenderer(Renderer):
-    def __init__(self):
-        self.names: dict[str, str] = {}
+    def __init__(self) -> None:
+        self.nodes: dict[str, Node] = {}
 
     def _render_node(self, node: Node) -> str:
+        if node.type is NodeType.ALIAS:
+            return ""
         name = self._node_type(node) + self._escape(node.name)
         return (
             f'class {node.id}["{name}"]'
             + "{\n"
             + ((f"    {node.source_name}()\n") if node.source_name else " \n")
             + "".join(
-                f"    {self.names[dep]}\n" for dep in node.dependencies
+                f"    {self.nodes[dep].name}\n" for dep in node.dependencies
             )
             + "}\n"
         )
@@ -44,17 +46,29 @@ class MermaidRenderer(Renderer):
         return re.sub(r"[^\w_.\-]", "_", line)
 
     def _render_node_deps(self, node: Node) -> list[str]:
-        return [
-            f"{dep} ..> {node.id}"
-            for dep in node.dependencies
-        ]
+        res: list[str] = []
+        if node.type is NodeType.ALIAS:
+            return res
+        for dep in node.dependencies:
+            dep_node = self.nodes[dep]
+            if dep_node.type is NodeType.ALIAS:
+                target = dep_node.dependencies[0]
+                res.append(
+                    f"{target} <.. {node.id}: "
+                    f"🔗 {self._escape(dep_node.name)}",
+                )
+            else:
+                res.append(f"{dep} <.. {node.id}")
+        return res
 
     def _node_type(self, node: Node) -> str:
         if node.is_protocol:
             prefix = "Ⓘ "
         else:
             prefix = ""
-        if node.type is NodeType.CONTEXT:
+        if node.type is NodeType.DECORATOR:
+            return "🎭 " + prefix
+        elif node.type is NodeType.CONTEXT:
             return "📥 " + prefix
         elif node.type is NodeType.ALIAS:
             return "🔗 " + prefix
@@ -85,14 +99,14 @@ class MermaidRenderer(Renderer):
             res += self._render_links(child)
         return res
 
-    def _fill_names(self, groups: list[Group]) -> None:
+    def _fill_nodes(self, groups: list[Group]) -> None:
         for group in groups:
             for node in group.nodes:
-                self.names[node.id] = node.name
-            self._fill_names(group.children)
+                self.nodes[node.id] = node
+            self._fill_nodes(group.children)
 
     def render(self, groups: list[Group]) -> str:
-        self._fill_names(groups)
+        self._fill_nodes(groups)
 
         res = "classDiagram\n"
         res += "direction LR\n"
diff --git a/src/dishka/plotter/model.py b/src/dishka/plotter/model.py
index add3ed8..992308e 100644
--- a/src/dishka/plotter/model.py
+++ b/src/dishka/plotter/model.py
@@ -11,6 +11,7 @@ class GroupType(Enum):
 class NodeType(Enum):
     CONTEXT = "Context"
     FACTORY = "Factory"
+    DECORATOR = "Decorator"
     ALIAS = "Alias"
 
 
diff --git a/src/dishka/plotter/transform.py b/src/dishka/plotter/transform.py
index ca90fed..9a4012a 100644
--- a/src/dishka/plotter/transform.py
+++ b/src/dishka/plotter/transform.py
@@ -4,13 +4,13 @@ from dishka import AsyncContainer, BaseScope, Container, DependencyKey
 from dishka._adaptix.type_tools import is_protocol
 from dishka.dependency_source import Factory
 from dishka.entities.factory_type import FactoryType
-from dishka.registry import Registry
+from dishka.registry import DECORATED_COMPONENT_PREFIX, Registry
 from dishka.text_rendering import get_name
 from .model import Group, GroupType, Node, NodeType
 
 
 class Transformer:
-    def __init__(self):
+    def __init__(self) -> None:
         self.nodes: dict[tuple[DependencyKey, BaseScope], Node] = {}
         self.groups: dict[Any, Group] = {}
         self._counter = 0
@@ -25,7 +25,18 @@ class Transformer:
                 return False
         return True
 
+    def _is_decorated(self, key: DependencyKey) -> bool:
+        if not key.component:
+            return False
+        return key.component.startswith(DECORATED_COMPONENT_PREFIX)
     def _node_type(self, factory: Factory) -> NodeType:
+        for dep in factory.dependencies:
+            if self._is_decorated(dep):
+                return NodeType.DECORATOR
+        for dep in factory.kw_dependencies.values():
+            if self._is_decorated(dep):
+                return NodeType.DECORATOR
+
         if factory.type is FactoryType.ALIAS:
             return NodeType.ALIAS
         elif factory.type is FactoryType.CONTEXT:
@@ -33,6 +44,28 @@ class Transformer:
         else:
             return NodeType.FACTORY
 
+    def _is_decorated_component(self, group: Group) -> bool:
+        if group.type is not GroupType.COMPONENT:
+            return False
+        return group.name.startswith(DECORATED_COMPONENT_PREFIX)
+
+    def _trace_decorator(self, node: Node, target_group: Group) -> None:
+        if node.type is not NodeType.DECORATOR:
+            return
+        for group in self.groups.values():
+            if not self._is_decorated_component(group):
+                continue
+            nodes_to_move = [
+                n
+                for n in group.nodes
+                if n.id in node.dependencies
+            ]
+            for moved in nodes_to_move:
+                group.nodes.remove(moved)
+                target_group.nodes.append(moved)
+                self._trace_decorator(moved, target_group)
+
+
     def _make_factories(
             self, scope: BaseScope, group: Group, registry: Registry,
     ) -> None:
@@ -50,13 +83,11 @@ class Transformer:
                 )
                 group.children.append(component_group)
             node_name = get_name(key.type_hint, include_module=False)
-            if key.component:
-                node_name += " " + str(key.component)
-
             if factory.type in (FactoryType.CONTEXT, FactoryType.ALIAS):
                 source_name = ""
             else:
                 source_name = get_name(factory.source, include_module=False)
+
             node = Node(
                 id=self.count("factory"),
                 name=node_name,
@@ -87,6 +118,14 @@ class Transformer:
                 dep_node = self.nodes[dep, dep_registry.scope]
                 node.dependencies.append(dep_node.id)
 
+    def clean_groups(self, groups: list[Group]) -> list[Group]:
+        res = []
+        for group in groups:
+            group.children = self.clean_groups(group.children)
+            if group.children or group.nodes:
+                res.append(group)
+        return res
+
     def transform(self, container: Container|AsyncContainer) -> list[Group]:
         registries = [container.registry, *container.child_registries]
         result = []
@@ -108,4 +147,9 @@ class Transformer:
             if self._is_empty(registry):
                 continue
             self._fill_dependencies(registry, registries[n::-1])
-        return result
+
+        for group in self.groups.values():
+            for node in group.nodes:
+                self._trace_decorator(node, group)
+
+        return self.clean_groups(result)
diff --git a/src/dishka/registry.py b/src/dishka/registry.py
index f699b00..e8915f3 100644
--- a/src/dishka/registry.py
+++ b/src/dishka/registry.py
@@ -29,6 +29,7 @@ from .exceptions import (
 from .factory_compiler import compile_factory
 from .provider import BaseProvider
 
+DECORATED_COMPONENT_PREFIX = "__Dishka_decorate_"
 
 class Registry:
     __slots__ = ("scope", "factories", "compiled", "compiled_async")
@@ -464,7 +465,7 @@ class RegistryBuilder:
         ):
             return
         depth = self.decorator_depth[provides]
-        decorated_component = f"__Dishka_decorate_{depth}"
+        decorated_component = f"{DECORATED_COMPONENT_PREFIX}{depth}"
         self.decorator_depth[provides] += 1
         if old_factory is None:
             raise InvalidGraphError(


Test patch:
diff --git a/tests/unit/plotter/test_transform.py b/tests/unit/plotter/test_transform.py
index b61d23f..d28b7c1 100644
--- a/tests/unit/plotter/test_transform.py
+++ b/tests/unit/plotter/test_transform.py
@@ -61,10 +61,10 @@ EXPECTED_GRAPH = [
                     Node(
                         id="factory6",
                         name="P",
-                        dependencies=[],
-                        type=NodeType.FACTORY,
+                        dependencies=["factory9"],
+                        type=NodeType.DECORATOR,
                         is_protocol=True,
-                        source_name="test_deps.<locals>.<lambda>",
+                        source_name="p_decorator",
                     ),
                     Node(
                         id="factory7",
@@ -74,6 +74,14 @@ EXPECTED_GRAPH = [
                         is_protocol=False,
                         source_name="",
                     ),
+                    Node(
+                        id="factory9",
+                        name="P",
+                        dependencies=[],
+                        type=NodeType.FACTORY,
+                        is_protocol=True,
+                        source_name="test_deps.<locals>.<lambda>",
+                    ),
                 ],
                 type=GroupType.COMPONENT,
             ),
@@ -84,11 +92,15 @@ EXPECTED_GRAPH = [
 ]
 
 
+def p_decorator(p: P) -> P:
+    return p
+
 def test_deps():
     provider = Provider(scope=Scope.APP)
     provider.provide(A)
     provider.provide(B)
     provider.provide(lambda: 0, provides=P)
+    provider.decorate(p_decorator)
     provider.alias(source=P, provides=float)
     container = make_container(provider)
     res = Transformer().transform(container)
@@ -215,7 +227,7 @@ COMPONENT_GRAPH = [
                 nodes=[
                     Node(
                         id="factory5",
-                        name="Container comp1",
+                        name="Container",
                         dependencies=["factory3"],
                         type=NodeType.ALIAS,
                         is_protocol=False,
@@ -223,7 +235,7 @@ COMPONENT_GRAPH = [
                     ),
                     Node(
                         id="factory6",
-                        name="A comp1",
+                        name="A",
                         dependencies=[],
                         type=NodeType.FACTORY,
                         is_protocol=False,