From 7a4a6afa23d873b797ba38498f2bb558b8eb4423 Mon Sep 17 00:00:00 2001
From: huyenngn <ngan-huyen.thi-nguyen@deutschebahn.com>
Date: Tue, 26 Mar 2024 17:44:54 +0100
Subject: [PATCH] feat: Add option to specify import package

---
 capella_ros_tools/__main__.py | 48 ++++++++++++++++++++++++++++++-----
 capella_ros_tools/importer.py | 39 +++++++++++++++++++++-------
 2 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/capella_ros_tools/__main__.py b/capella_ros_tools/__main__.py
index e655c8c..bc929e7 100644
--- a/capella_ros_tools/__main__.py
+++ b/capella_ros_tools/__main__.py
@@ -4,6 +4,7 @@
 
 import io
 import pathlib
+import uuid
 
 import capellambse
 import click
@@ -44,9 +45,20 @@ def cli():
     "-l",
     "--layer",
     type=click.Choice(["oa", "la", "sa", "pa"], case_sensitive=False),
-    required=True,
     help="The layer to import the messages to.",
 )
+@click.option(
+    "-r",
+    "--root",
+    type=click.UUID,
+    help="The UUID of the root package to import the messages to.",
+)
+@click.option(
+    "-t",
+    "--types",
+    type=click.UUID,
+    help="The UUID of the types package to import the created data types to.",
+)
 @click.option(
     "--no-deps",
     "no_deps",
@@ -63,18 +75,29 @@ def import_msgs(
     input: str,
     model: capellambse.MelodyModel,
     layer: str,
+    root: uuid.UUID,
+    types: uuid.UUID,
     no_deps: bool,
     output: pathlib.Path,
 ) -> None:
     """Import ROS messages into a Capella data package."""
 
-    root_uuid = getattr(model, layer).data_package.uuid
-    types_uuid = model.sa.data_package.uuid
+    if root:
+        root_uuid = str(root)
+    elif layer:
+        root_uuid = getattr(model, layer).data_package.uuid
+    else:
+        raise click.UsageError("Either --root or --layer must be provided")
+
+    if types:
+        params = {"types_uuid": str(types)}
+    else:
+        params = {"types_parent_uuid": model.sa.data_package.uuid}
 
     parsed = importer.Importer(input, no_deps)
     logger.info("Loaded %d packages", len(parsed.messages.packages))
 
-    yml = parsed.to_yaml(root_uuid, types_uuid)
+    yml = parsed.to_yaml(root_uuid, **params)
     if output:
         logger.info("Writing to file %s", output)
         output.write_text(yml, encoding="utf-8")
@@ -96,9 +119,14 @@ def import_msgs(
     "-l",
     "--layer",
     type=click.Choice(["oa", "la", "sa", "pa"], case_sensitive=False),
-    required=True,
     help="The layer to export the model objects from.",
 )
+@click.option(
+    "-r",
+    "--root",
+    type=click.UUID,
+    help="The UUID of the root package to import the messages to.",
+)
 @click.option(
     "-o",
     "--output",
@@ -109,11 +137,17 @@ def import_msgs(
 def export_capella(
     model: capellambse.MelodyModel,
     layer: str,
+    root: uuid.UUID,
     output: pathlib.Path,
 ):
     """Export Capella data package to ROS messages."""
-    current_pkg = getattr(model, layer).data_package
-    exporter.export(current_pkg, output)
+    if root:
+        current_pkg = model.search("DataPkg").by_uuid(str(root))
+    elif layer:
+        current_pkg = getattr(model, layer).data_package
+    else:
+        raise click.UsageError("Either --root or --layer must be provided")
+    exporter.export(current_pkg, output)  # type: ignore
 
 
 if __name__ == "__main__":
diff --git a/capella_ros_tools/importer.py b/capella_ros_tools/importer.py
index 85a1868..e7e9afe 100644
--- a/capella_ros_tools/importer.py
+++ b/capella_ros_tools/importer.py
@@ -196,24 +196,41 @@ def _convert_enum(
 
         return yml
 
-    def to_yaml(self, layer_data_uuid: str, sa_data_uuid: str) -> str:
+    def to_yaml(
+        self,
+        root_uuid: str,
+        types_parent_uuid: str = "",
+        types_uuid: str = "",
+    ) -> str:
         """Import ROS messages into a Capella data package."""
         logger.info("Generating decl YAML")
         instructions = [
-            {"parent": decl.UUIDReference(helpers.UUIDString(layer_data_uuid))}
+            {"parent": decl.UUIDReference(helpers.UUIDString(root_uuid))}
             | self._convert_package(self.messages),
         ]
-        if needed_types := [
+        needed_types = [
             p for p in self._promise_id_refs if p not in self._promise_ids
-        ]:
-            datatypes = [
-                self._convert_datatype(promise_id)
-                for promise_id in needed_types
-            ]
+        ]
+        if not needed_types:
+            return decl.dump(instructions)
+
+        datatypes = [
+            self._convert_datatype(promise_id) for promise_id in needed_types
+        ]
+        if types_uuid:
+            instructions.append(
+                {
+                    "parent": decl.UUIDReference(
+                        helpers.UUIDString(types_uuid)
+                    ),
+                    "sync": {"datatypes": datatypes},
+                }
+            )
+        elif types_parent_uuid:
             instructions.append(
                 {
                     "parent": decl.UUIDReference(
-                        helpers.UUIDString(sa_data_uuid)
+                        helpers.UUIDString(types_parent_uuid)
                     ),
                     "sync": {
                         "packages": [
@@ -225,4 +242,8 @@ def to_yaml(self, layer_data_uuid: str, sa_data_uuid: str) -> str:
                     },
                 }
             )
+        else:
+            raise ValueError(
+                "Either types_parent_uuid or types_uuid must be provided"
+            )
         return decl.dump(instructions)