From fd9bc9c3bc4e0f6fe4b2d55fc9166202912ad7b5 Mon Sep 17 00:00:00 2001 From: Sergei Konik Date: Sun, 22 Oct 2023 22:00:12 +0300 Subject: [PATCH] feat: add autoimport of all diagrams classes --- .flake8 | 2 +- docker_compose_diagram/constants.py | 2 + .../docker_images/auto_import_diagrams.py | 75 ------------- .../docker_images/diagrams_nodes_importer.py | 105 ++++++++++++++++++ .../docker_images/patterns.py | 8 +- docker_compose_diagram/docker_images/utils.py | 22 +++- pyproject.toml | 2 +- 7 files changed, 132 insertions(+), 84 deletions(-) delete mode 100644 docker_compose_diagram/docker_images/auto_import_diagrams.py create mode 100644 docker_compose_diagram/docker_images/diagrams_nodes_importer.py diff --git a/.flake8 b/.flake8 index d0e140c..254c92b 100644 --- a/.flake8 +++ b/.flake8 @@ -1,6 +1,6 @@ [flake8] max-line-length=100 -max-complexity=5 +max-complexity=6 per-file-ignores = __init__.py: F401, W605 diff --git a/docker_compose_diagram/constants.py b/docker_compose_diagram/constants.py index d06afdb..f817b1d 100644 --- a/docker_compose_diagram/constants.py +++ b/docker_compose_diagram/constants.py @@ -2,3 +2,5 @@ IMAGE_EXTENSION_RE = re.compile(r".*\.(png)$") +PACKAGE_CLASS_PATH = re.compile(r"^(?P[^:]*)\.(?P[A-z].*)") +DIAGRAMS_PACKAGE_NAME = "diagrams" diff --git a/docker_compose_diagram/docker_images/auto_import_diagrams.py b/docker_compose_diagram/docker_images/auto_import_diagrams.py deleted file mode 100644 index 59da01f..0000000 --- a/docker_compose_diagram/docker_images/auto_import_diagrams.py +++ /dev/null @@ -1,75 +0,0 @@ -from typing import Any, Type - -from diagrams import Node -from diagrams.onprem import _OnPrem - - -def import_on_prem_parent_nodes() -> list[Type[_OnPrem]]: - from diagrams.onprem.aggregator import _Aggregator - from diagrams.onprem.analytics import _Analytics - from diagrams.onprem.auth import _Auth - from diagrams.onprem.cd import _Cd - from diagrams.onprem.certificates import _Certificates - from diagrams.onprem.ci import _Ci - from diagrams.onprem.client import _Client - from diagrams.onprem.compute import _Compute - from diagrams.onprem.database import _Database - from diagrams.onprem.dns import _Dns - from diagrams.onprem.etl import _Etl - from diagrams.onprem.gitops import _Gitops - from diagrams.onprem.inmemory import _Inmemory - from diagrams.onprem.logging import _Logging - from diagrams.onprem.network import _Network - from diagrams.onprem.queue import _Queue - from diagrams.onprem.storage import _Storage - - return [ - _Database, - _Queue, - _Network, - _Inmemory, - _Storage, - _Logging, - _Auth, - _Client, - _Analytics, - _Aggregator, - _Certificates, - _Cd, - _Ci, - _Compute, - _Dns, - _Etl, - _Gitops, - ] - - -def _collect_subclasses( - diagrams_parent_class: Type[Node], base_class: Type[Any] -) -> list[Type[Any]]: - new_classes = [] - for diagram_child_class in diagrams_parent_class.__subclasses__(): - new_class = type( - f"{diagram_child_class.__name__}Image", - (base_class,), - { - "pattern": str(diagram_child_class.__name__).lower(), - "diagram_render_class": diagram_child_class, - }, - ) - new_classes.append(new_class) - - return new_classes - - -def register_all_icons_from_diagrams(base_class: Type[Any]) -> list[Type[Any]]: - """Collect""" - new_classes = [] - - for on_prem_diagrams_parent in import_on_prem_parent_nodes(): - database_classes = _collect_subclasses( - diagrams_parent_class=on_prem_diagrams_parent, base_class=base_class - ) - new_classes.extend(database_classes) - - return new_classes diff --git a/docker_compose_diagram/docker_images/diagrams_nodes_importer.py b/docker_compose_diagram/docker_images/diagrams_nodes_importer.py new file mode 100644 index 0000000..e8cbf2f --- /dev/null +++ b/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 diff --git a/docker_compose_diagram/docker_images/patterns.py b/docker_compose_diagram/docker_images/patterns.py index d6c3dab..1b81a37 100644 --- a/docker_compose_diagram/docker_images/patterns.py +++ b/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 +) diff --git a/docker_compose_diagram/docker_images/utils.py b/docker_compose_diagram/docker_images/utils.py index 7a661bb..87a4b05 100644 --- a/docker_compose_diagram/docker_images/utils.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 230025f..650b6e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "docker-compose-diagram" -version = "0.5.0" +version = "0.6.0" description = "" authors = ["Sergei Konik "]