From 84b24bcb0bf977ce788ec9be70c81a5e117061fb Mon Sep 17 00:00:00 2001 From: Keaton Cross Date: Thu, 22 Jun 2023 12:34:17 -0500 Subject: [PATCH] fix: unify generation code for consistent imports --- pystructurizr/cli.py | 17 +++++++---------- pystructurizr/cli_helper.py | 25 +++++-------------------- pystructurizr/generator.py | 37 +++++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 30 deletions(-) create mode 100644 pystructurizr/generator.py diff --git a/pystructurizr/cli.py b/pystructurizr/cli.py index ac5e24a..321b41b 100644 --- a/pystructurizr/cli.py +++ b/pystructurizr/cli.py @@ -6,8 +6,7 @@ import click -from .cli_helper import (ensure_tmp_folder_exists, generate_diagram_code, - generate_diagram_code_in_child_process, generate_svg) +from .cli_helper import (ensure_tmp_folder_exists, generate_diagram_code_in_child_process, generate_svg) from .cli_watcher import observe_modules from .cloudstorage import CloudStorage, create_cloud_storage @@ -18,7 +17,7 @@ @click.option('--as-json', is_flag=True, default=False, help='Dumps the generated code and the imported modules as a json object') def dump(view, as_json): - diagram_code, imported_modules = generate_diagram_code(view) + diagram_code, imported_modules = generate_diagram_code_in_child_process(view) if as_json: print(json.dumps({ "code": diagram_code, @@ -28,7 +27,6 @@ def dump(view, as_json): print(diagram_code) - @click.command() @click.option('--view', prompt='Your view file (e.g. example.componentview)', help='The view file to develop.') @@ -42,9 +40,9 @@ def dev(view): async def async_behavior(): print("Generating diagram...") - result = generate_diagram_code_in_child_process(view) - await generate_svg(result['code'], tmp_folder) - return result['imported_modules'] + diagram_code, imported_modules = generate_diagram_code_in_child_process(view) + await generate_svg(diagram_code, tmp_folder) + return imported_modules async def observe_loop(): modules_to_watch = await async_behavior() @@ -56,7 +54,6 @@ async def observe_loop(): asyncio.run(observe_loop()) - @click.command() @click.option('--view', prompt='Your view file (e.g. example.componentview)', help='The view file to generate and upload to cloud storage.') @@ -69,7 +66,7 @@ async def observe_loop(): def build(view, gcs_credentials, bucket_name, object_name): async def async_behavior(): # Generate diagram - diagram_code, _ = generate_diagram_code(view) + diagram_code, _ = generate_diagram_code_in_child_process(view) tmp_folder = ensure_tmp_folder_exists() # Generate SVG @@ -83,11 +80,11 @@ async def async_behavior(): asyncio.run(async_behavior()) - @click.group() def cli(): pass + cli.add_command(dump) cli.add_command(dev) cli.add_command(build) diff --git a/pystructurizr/cli_helper.py b/pystructurizr/cli_helper.py index c7b5217..89ffcad 100644 --- a/pystructurizr/cli_helper.py +++ b/pystructurizr/cli_helper.py @@ -1,4 +1,3 @@ -import importlib import json import os import subprocess @@ -9,33 +8,19 @@ import httpx -def generate_diagram_code(view: str) -> str: - try: - initial_modules = set(sys.modules.keys()) - module = importlib.import_module(view) - imported_modules = set(sys.modules.keys()) - initial_modules - code = module.workspace.dump() - return code, imported_modules - except ModuleNotFoundError: - # pylint: disable=raise-missing-from - raise click.BadParameter("Invalid view name. Make sure you don't include the .py file extension.") - except AttributeError: - # pylint: disable=raise-missing-from - raise click.BadParameter("Non-compliant view file: make sure it exports the PyStructurizr workspace.") - - -def generate_diagram_code_in_child_process(view: str) -> str: +def generate_diagram_code_in_child_process(view: str) -> tuple[dict, list[str]]: def run_child_process(): # Run a separate Python script as a child process - output = subprocess.check_output([sys.executable, "-m", "pystructurizr.cli", "dump", "--as-json", "--view", view]) + output = subprocess.check_output([sys.executable, "-m", "pystructurizr.generator", "dump", "--view", view]) return output.decode().strip() # Run the child process and capture its output child_output = run_child_process() - return json.loads(child_output) + result = json.loads(child_output) + return result['code'], result['imported_modules'] -async def generate_svg(diagram_code: str, tmp_folder: str) -> str: +async def generate_svg(diagram_code: dict, tmp_folder: str) -> str: url = "https://kroki.io/structurizr/svg" async with httpx.AsyncClient() as client: resp = await client.post(url, data=diagram_code) diff --git a/pystructurizr/generator.py b/pystructurizr/generator.py new file mode 100644 index 0000000..f9e633c --- /dev/null +++ b/pystructurizr/generator.py @@ -0,0 +1,37 @@ +import importlib +import json +import sys + +import click + + +@click.command() +@click.option('--view', prompt='Your view file (e.g. example.componentview)', + help='The view file to generate.') +def dump(view: str): + try: + initial_modules = set(sys.modules.keys()) + module = importlib.import_module(view) + imported_modules = set(sys.modules.keys()) - initial_modules + code = module.workspace.dump() + print(json.dumps({ + "code": code, + "imported_modules": list(imported_modules) + })) + except ModuleNotFoundError: + # pylint: disable=raise-missing-from + raise click.BadParameter("Invalid view name. Make sure you don't include the .py file extension.") + except AttributeError: + # pylint: disable=raise-missing-from + raise click.BadParameter("Non-compliant view file: make sure it exports the PyStructurizr workspace.") + + +@click.group() +def cli(): + pass + + +cli.add_command(dump) + +if __name__ == '__main__': + cli()