Merge pull request #13 from skonik/feature/finish-automatic-diagrams-…

…imports

feat: add autoimport of all diagrams classes

.flake8

@@ -1,6 +1,6 @@
 [flake8]
 max-line-length=100
-max-complexity=5
+max-complexity=6
 
 per-file-ignores =
     __init__.py: F401, W605

docker_compose_diagram/constants.py

@@ -2,3 +2,5 @@
 
 
 IMAGE_EXTENSION_RE = re.compile(r".*\.(png)$")
+PACKAGE_CLASS_PATH = re.compile(r"^(?P<package_path>[^:]*)\.(?P<class_name>[A-z].*)")
+DIAGRAMS_PACKAGE_NAME = "diagrams"

docker_compose_diagram/docker_images/diagrams_nodes_importer.py

@@ -0,0 +1,105 @@
+import importlib
+import inspect
+import pkgutil
+from types import ModuleType
+from typing import Any, Type
+
+from diagrams import Node
+
+from docker_compose_diagram.constants import DIAGRAMS_PACKAGE_NAME
+
+
+FIRST_VARIABLE_INDEX = 0
+FIRST_BASE_CLASS_INDEX = 0
+PARENT_CLASS_PREFIX = "_"
+
+PATTERNS_REFORMAT = {"r": r"^R$"}
+
+
+def import_all_parent_nodes():
+    """
+    Diagrams packages has parent nodes like _Database
+    which are inherited by actual icon nodes.
+    We can search for all such subclasses.
+
+    Firstly we can iterate over all packages inside diagrams nodes.
+    Then iterate over each file in the package.
+    And find classes that start with `_`.
+    We will consider them as parent classes for our nodes.
+    Then we can just return these parent nodes.
+    """
+    source_package_module: ModuleType = importlib.import_module(DIAGRAMS_PACKAGE_NAME)
+    source_package_module_path = source_package_module.__path__
+
+    all_parent_nodes = []
+    # high_level_module is something like diagrams.aws, diagrams.onprem, etc
+    for high_level_module in pkgutil.iter_modules(source_package_module_path):
+
+        diagrams_high_level_package_name = ".".join(
+            [DIAGRAMS_PACKAGE_NAME, high_level_module.name]
+        )
+        diagrams_high_level_package: ModuleType = importlib.import_module(
+            diagrams_high_level_package_name
+        )
+        diagrams_high_level_package_path = diagrams_high_level_package.__path__
+
+        # low level module is something like diagrams.onprem.database, etc
+        for low_level_module in pkgutil.iter_modules(diagrams_high_level_package_path):
+            diagrams_low_level_module_name = ".".join(
+                [diagrams_high_level_package_name, low_level_module.name]
+            )
+            diagrams_low_level_module: ModuleType = importlib.import_module(
+                diagrams_low_level_module_name
+            )
+
+            module_variables = inspect.getmembers(diagrams_low_level_module)
+
+            name, obj = module_variables[FIRST_VARIABLE_INDEX]
+            if name.startswith(PARENT_CLASS_PREFIX):
+                all_parent_nodes.append(obj)
+                continue
+
+            first_parent_class = obj.__bases__[FIRST_BASE_CLASS_INDEX]
+            all_parent_nodes.append(first_parent_class)
+
+    return all_parent_nodes
+
+
+def _create_docker_image_classes(
+    diagrams_parent_class: Type[Node], base_class: Type[Any]
+) -> list[Type[Any]]:
+    new_classes = []
+    for diagram_child_class in diagrams_parent_class.__subclasses__():
+        pattern = str(diagram_child_class.__name__).lower()
+        if pattern in PATTERNS_REFORMAT:
+            pattern = PATTERNS_REFORMAT[pattern]
+
+        new_class = type(
+            f"{diagram_child_class.__name__}Image",
+            (base_class,),
+            {
+                "pattern": pattern,
+                "diagram_render_class": diagram_child_class,
+            },
+        )
+        new_classes.append(new_class)
+
+    return new_classes
+
+
+def create_docker_image_patterns_from_diagrams(
+    base_class: Type[Any],
+) -> list[Type[Any]]:
+    """
+    Collect all diagrams nodes and create proper wrappers
+    (classes which inherit our DockerImagePattern class)
+    """
+    all_nodes = []
+
+    for parent_node in import_all_parent_nodes():
+        nodes = _create_docker_image_classes(
+            diagrams_parent_class=parent_node, base_class=base_class
+        )
+        all_nodes.extend(nodes)
+
+    return all_nodes

docker_compose_diagram/docker_images/patterns.py

@@ -54,8 +54,8 @@
     Typescript,
 )
 
-from docker_compose_diagram.docker_images.auto_import_diagrams import (
-    register_all_icons_from_diagrams,
+from docker_compose_diagram.docker_images.diagrams_nodes_importer import (
+    create_docker_image_patterns_from_diagrams,
 )
 
 
@@ -360,4 +360,6 @@ class LoadBalancerImage(DockerImagePattern):
     diagram_render_class = LoadBalancer
 
 
-diagrams_classes = register_all_icons_from_diagrams(base_class=DockerImagePattern)
+diagrams_classes = create_docker_image_patterns_from_diagrams(
+    base_class=DockerImagePattern
+)

docker_compose_diagram/docker_images/utils.py

@@ -1,3 +1,4 @@
+import importlib
 import re
 from os import path
 from typing import Any, Dict, Optional, Type, Union
@@ -7,7 +8,7 @@
 from diagrams.generic.compute import Rack
 from dockerfile_parse import DockerfileParser
 
-from docker_compose_diagram.constants import IMAGE_EXTENSION_RE
+from docker_compose_diagram.constants import IMAGE_EXTENSION_RE, PACKAGE_CLASS_PATH
 from docker_compose_diagram.docker_images.patterns import DockerImagePattern
 
 
@@ -47,18 +48,31 @@ def determine_image_name(
     return image_name
 
 
+def _import_node_class_from_path(image_name: str) -> Type[Node]:
+    search_result = PACKAGE_CLASS_PATH.search(image_name)
+    class_name = search_result.groupdict()["class_name"]
+    package_path = search_result.groupdict()["package_path"]
+    module = importlib.import_module(package_path)
+    return getattr(module, class_name)
+
+
 def determine_diagram_render_class(
     image_name: str,
 ) -> Union[Type[DockerImagePattern], Type[Node]]:
     if image_name is None:
         return DEFAULT_ICON_CLASS
-
     # in case if user provides "example.svg" as icon name
-    if IMAGE_EXTENSION_RE.match(image_name):
+    elif IMAGE_EXTENSION_RE.match(image_name):
         return Custom
+    # If we provide path to class
+    # diagrams.onprem.analytics.Beam
+    elif PACKAGE_CLASS_PATH.match(image_name):
+        return _import_node_class_from_path(image_name)
 
-    for subclass in DockerImagePattern.__subclasses__():
+    subclasses = DockerImagePattern.__subclasses__()
+    for subclass in subclasses:
         re_match = re.search(subclass.pattern, image_name)
+
         if re_match:
             return subclass.diagram_render_class
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "docker-compose-diagram"
-version = "0.5.0"
+version = "0.6.0"
 description = ""
 authors = ["Sergei Konik <s.konik.dev@gmail.com>"]