Refactor/improve tes models (#24)

Co-authored-by: salihuDickson <[email protected]>

.github/workflows/ci.yml

@@ -2,43 +2,44 @@ name: CI
 
 on:
   push:
-    branches: [ main ]
+    branches: '*'
   pull_request:
-    branches: [ main ]
+    branches: '*'
 
 jobs:
   build:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
-
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.11'
-
-    - name: Install Poetry
-      run: |
-        curl -sSL https://install.python-poetry.org | python3 -
-        poetry install
-
-    - name: Lint with Ruff
-      run: |
-        poetry run ruff check crategen/
-
-    - name: Type check with Mypy
-      run: |
-        poetry run mypy crategen/
-
-    - name: Run security checks with Bandit
-      run: |
-        poetry run bandit -r crategen/
-
-    - name: Install test dependencies
-      run: |
-        poetry add pytest pytest-cov pytest-mock
-
-    - name: Run tests
-      run: |
-        poetry run pytest --cov=crategen
+      - uses: actions/checkout@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.11'
+
+      - name: Install Poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+          poetry install
+
+      - name: Lint with Ruff
+        run: |
+          poetry run ruff check crategen/
+        if: ${{ success() }}
+
+      - name: Type check with Mypy
+        run: |
+          poetry run mypy crategen/
+
+      - name: Run security checks with Bandit
+        run: |
+          poetry run bandit -r crategen/
+
+      - name: Install test dependencies
+        run: |
+          poetry add pytest pytest-cov pytest-mock
+
+      - name: Run tests
+        run: |
+          poetry run pytest --cov=crategen

crategen/cli.py

@@ -1,30 +1,37 @@
-import click
 import json
+
+import click
+
 from crategen.converter_manager import ConverterManager
 
+
 @click.command()
-@click.option('--input', prompt='Input file', help='Path to the input JSON file.')
-@click.option('--output', prompt='Output file', help='Path to the output JSON file.')
-@click.option('--conversion-type', prompt='Conversion type', type=click.Choice(['tes-to-wrroc', 'wes-to-wrroc']), help='Type of conversion to perform.')
+@click.option("--input", prompt="Input file", help="Path to the input JSON file.")
+@click.option("--output", prompt="Output file", help="Path to the output JSON file.")
+@click.option(
+    "--conversion-type",
+    prompt="Conversion type",
+    type=click.Choice(["tes-to-wrroc", "wes-to-wrroc"]),
+    help="Type of conversion to perform.",
+)
 def cli(input, output, conversion_type):
-    """
-    Command Line Interface for converting TES/WES to WRROC.
-    """
+    """Command Line Interface for converting TES/WES to WRROC."""
     manager = ConverterManager()
 
     # Load input data from JSON file
-    with open(input, 'r') as input_file:
+    with open(input) as input_file:
         data = json.load(input_file)
 
     # Perform the conversion based on the specified type
-    if conversion_type == 'tes-to-wrroc':
+    if conversion_type == "tes-to-wrroc":
         result = manager.convert_tes_to_wrroc(data)
-    elif conversion_type == 'wes-to-wrroc':
+    elif conversion_type == "wes-to-wrroc":
         result = manager.convert_wes_to_wrroc(data)
-    
+
     # Save the result to the output JSON file
-    with open(output, 'w') as output_file:
+    with open(output, "w") as output_file:
         json.dump(result, output_file, indent=4)
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     cli()

crategen/converter_manager.py

@@ -1,6 +1,7 @@
 from .converters.tes_converter import TESConverter
 from .converters.wes_converter import WESConverter
 
+
 class ConverterManager:
     def __init__(self):
         self.tes_converter = TESConverter()

crategen/converters/abstract_converter.py

@@ -1,10 +1,11 @@
 from abc import ABC, abstractmethod
 
+
 class AbstractConverter(ABC):
     @abstractmethod
     def convert_to_wrroc(self, data):
         """Convert data to WRROC format"""
-    
+
     @abstractmethod
     def convert_from_wrroc(self, wrroc_data):
         """Convert WRROC data to the original format"""

crategen/converters/tes_converter.py

@@ -1,8 +1,10 @@
-from .abstract_converter import AbstractConverter
-from .utils import convert_to_iso8601
-from ..models import TESData, WRROCDataTES
 from pydantic import ValidationError
 
+from ..models.tes_models import TESData
+from ..models.wrroc_models import WRROCDataTES
+from .abstract_converter import AbstractConverter
+
+
 class TESConverter(AbstractConverter):
     def convert_to_wrroc(self, data: dict) -> dict:
         """
@@ -23,16 +25,38 @@ def convert_to_wrroc(self, data: dict) -> dict:
         except ValidationError as e:
             raise ValueError(f"Invalid TES data: {e.errors()}") from e
 
+        # Extract validated data
+        (
+            id,
+            name,
+            description,
+            creation_time,
+            state,
+            inputs,
+            outputs,
+            executors,
+            resources,
+            volumes,
+            logs,
+            tags,
+        ) = data_tes.dict().values()
+        end_time = logs[0].end_time
+
         # Convert to WRROC format
         wrroc_data = {
-            "@id": data_tes.id,
-            "name": data_tes.name,
-            "description": data_tes.description,
-            "instrument": data_tes.executors[0].image if data_tes.executors else None,
-            "object": [{"@id": input.url, "name": input.path} for input in data_tes.inputs],
-            "result": [{"@id": output.url, "name": output.path} for output in data_tes.outputs],
-            "startTime": convert_to_iso8601(data_tes.creation_time),
-            "endTime": convert_to_iso8601(data_tes.logs[0].end_time if data_tes.logs else ""),
+            "@id": id,
+            "name": name,
+            "description": description,
+            "instrument": executors[0]["image"] if executors else None,
+            "object": [
+                {"@id": input["url"], "name": input["path"], "type": input["type"]}
+                for input in inputs
+            ],
+            "result": [
+                {"@id": output["url"], "name": output["path"]} for output in outputs
+            ],
+            "startTime": creation_time,
+            "endTime": end_time,
         }
         return wrroc_data
 

crategen/converters/wes_converter.py

@@ -1,8 +1,11 @@
-from .abstract_converter import AbstractConverter
-from .utils import convert_to_iso8601
-from ..models import WESData, WRROCDataWES
 from pydantic import ValidationError
 
+from ..models.wes_models import WESData
+from ..models.wrroc_models import WRROCDataWES
+from ..utils import convert_to_iso8601
+from .abstract_converter import AbstractConverter
+
+
 class WESConverter(AbstractConverter):
     def convert_to_wrroc(self, data: dict) -> dict:
         """
@@ -30,7 +33,10 @@ def convert_to_wrroc(self, data: dict) -> dict:
             "status": data_wes.state,
             "startTime": convert_to_iso8601(data_wes.run_log.start_time),
             "endTime": convert_to_iso8601(data_wes.run_log.end_time),
-            "result": [{"@id": output.location, "name": output.name} for output in data_wes.outputs],
+            "result": [
+                {"@id": output.location, "name": output.name}
+                for output in data_wes.outputs
+            ],
         }
         return wrroc_data
 
@@ -51,7 +57,9 @@ def convert_from_wrroc(self, data: dict) -> dict:
         try:
             data_wrroc = WRROCDataWES(**data)
         except ValidationError as e:
-            raise ValueError(f"Invalid WRROC data for WES conversion: {e.errors()}") from e
+            raise ValueError(
+                f"Invalid WRROC data for WES conversion: {e.errors()}"
+            ) from e
 
         # Convert from WRROC to WES format
         wes_data = {
@@ -62,6 +70,8 @@ def convert_from_wrroc(self, data: dict) -> dict:
                 "end_time": data_wrroc.endTime,
             },
             "state": data_wrroc.status,
-            "outputs": [{"location": res.id, "name": res.name} for res in data_wrroc.result],
+            "outputs": [
+                {"location": res.id, "name": res.name} for res in data_wrroc.result
+            ],
         }
         return wes_data