Merge pull request #3 from crossk3/kcross/unify-generation-code

fix: unify generation code for consistent imports

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)

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)

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()