Add GEM Transform

fia_api/scripts/transforms/factory.py

@@ -3,6 +3,7 @@
 import logging
 
 from fia_api.scripts.transforms.enginx_transform import EnginxTransform
+from fia_api.scripts.transforms.gem_transform import GEMTransform
 from fia_api.scripts.transforms.imat_transforms import IMATTransform
 from fia_api.scripts.transforms.iris_transform import IrisTransform
 from fia_api.scripts.transforms.mari_transforms import MariTransform
@@ -16,7 +17,7 @@
 logger = logging.getLogger(__name__)
 
 
-def get_transform_for_instrument(instrument: str) -> Transform:  # noqa: PLR0911
+def get_transform_for_instrument(instrument: str) -> Transform:  # noqa: PLR0911, C901
     """
     Get the appropriate transform for the given instrument and run file
     :param instrument: str - the instrument
@@ -40,6 +41,8 @@ def get_transform_for_instrument(instrument: str) -> Transform:  # noqa: PLR0911
             return EnginxTransform()
         case "imat":
             return IMATTransform()
+        case "gem":
+            return GEMTransform()
         case "test":
             return TestTransform()
         case _:

fia_api/scripts/transforms/gem_transform.py

@@ -0,0 +1,62 @@
+import logging
+
+from fia_api.core.models import Job
+from fia_api.scripts.pre_script import PreScript
+from fia_api.scripts.transforms.transform import Transform
+
+logger = logging.getLogger(__name__)
+
+
+class GEMTransform(Transform):
+    """
+    GEMTransform applies modifications to GEM instrument scripts based on reduction input parameters in a Reduction
+    entity.
+    """
+
+    def apply(self, script: PreScript, job: Job) -> None:  # noqa: PLR0912,C901
+        logger.info("Beginning GEM transform for job %s...", job.id)
+        lines = script.value.splitlines()
+        # MyPY does not believe ColumnElement[JSONB] is indexable, despite JSONB implementing the Indexable mixin
+        # If you get here in the future, try removing the following line and see if it passes with newer mypy.
+
+        runno = job.inputs["runno"]  # type: ignore
+        if isinstance(runno, list):
+            if len(runno) > 1:
+                # Convert list to range string if contiguous, otherwise comma-separated
+                if all(runno[i] == runno[i - 1] + 1 for i in range(1, len(runno))):
+                    runno_str = f"{runno[0]}-{runno[-1]}"
+                else:
+                    runno_str = ",".join(map(str, runno))
+            else:
+                runno_str = str(runno[0])
+        else:
+            runno_str = str(runno)
+
+        for index, line in enumerate(lines):
+            if line.startswith("mode ="):
+                lines[index] = f'mode = "{job.inputs["mode"]}"'  # type: ignore
+                continue
+            if line.startswith("input_mode ="):
+                lines[index] = f'input_mode = "{job.inputs["input_mode"]}"'  # type: ignore
+                continue
+            if line.startswith("vanadium_runno ="):
+                lines[index] = f"vanadium_runno = {runno_str}"
+                continue
+            if line.startswith("runno ="):
+                lines[index] = f"runno = {runno_str}"
+                continue
+            if line.startswith("calibration_dir ="):
+                lines[index] = f"calibration_dir = {job.inputs['calibration_dir']}"  # type: ignore
+                continue
+            if line.startswith("splined_vanadium_dir ="):
+                lines[index] = f'splined_vanadium_dir = "{job.inputs["splined_vanadium_dir"]}"'  # type: ignore
+                continue
+            if line.startswith("config_file ="):
+                lines[index] = f'config_file = "{job.inputs["config_file"]}"'  # type: ignore
+                continue
+            if line.startswith("output_dir = "):
+                lines[index] = f'output_dir = "{job.inputs["output_dir"]}"'  # type: ignore
+                continue
+
+        script.value = "\n".join(lines)
+        logger.info("Transform complete for job %s", job.id)

test/scripts/transforms/test_factory.py

@@ -4,6 +4,7 @@
 
 from fia_api.scripts.transforms.enginx_transform import EnginxTransform
 from fia_api.scripts.transforms.factory import get_transform_for_instrument
+from fia_api.scripts.transforms.gem_transform import GEMTransform
 from fia_api.scripts.transforms.imat_transforms import IMATTransform
 from fia_api.scripts.transforms.iris_transform import IrisTransform
 from fia_api.scripts.transforms.mari_transforms import MariTransform
@@ -28,6 +29,7 @@
         ("vesuvio", VesuvioTransform),
         ("enginx", EnginxTransform),
         ("imat", IMATTransform),
+        ("gem", GEMTransform),
     ],
 )
 def test_transform_factory(name, expected_transform):

test/scripts/transforms/test_gem_transform.py

@@ -0,0 +1,121 @@
+"""Test cases for GEMTransform."""
+
+from unittest.mock import Mock
+
+import pytest
+
+from fia_api.scripts.pre_script import PreScript
+from fia_api.scripts.transforms.gem_transform import GEMTransform
+
+SCRIPT = """
+mode = "default_mode"
+input_mode = "default_input_mode"
+vanadium_runno = 0
+runno = 0
+calibration_dir = None
+splined_vanadium_dir = "default_splined_vanadium_dir"
+config_file = "default_config_file"
+output_dir = "default_output_dir"
+"""
+
+
+@pytest.fixture
+def base_job():
+    """Fixture for base job inputs."""
+    job = Mock()
+    job.id = "test-job-gem"
+    job.inputs = {
+        "mode": "transmission",
+        "input_mode": "raw",
+        "calibration_dir": "/path/to/cal",
+        "splined_vanadium_dir": "/path/to/splined",
+        "config_file": "/path/to/config",
+        "output_dir": "/path/to/output",
+        "runno": 12345,
+    }
+    return job
+
+
+@pytest.fixture
+def create_expected_script():
+    """Fixture returning a helper function to construct the expected script with runno."""
+
+    def _create(runno_str: str) -> str:
+        return f"""
+mode = "transmission"
+input_mode = "raw"
+vanadium_runno = {runno_str}
+runno = {runno_str}
+calibration_dir = /path/to/cal
+splined_vanadium_dir = "/path/to/splined"
+config_file = "/path/to/config"
+output_dir = "/path/to/output\""""
+
+    return _create
+
+
+def test_gem_transform_single_run(base_job, create_expected_script):
+    """Test GEMTransform with a single run number."""
+    script = PreScript(value=SCRIPT)
+    GEMTransform().apply(script, base_job)
+
+    assert script.value == create_expected_script("12345")
+
+
+def test_gem_transform_contiguous_runs(base_job, create_expected_script):
+    """Test GEMTransform with contiguous runs."""
+    base_job.inputs["runno"] = [12345, 12346, 12347]
+    script = PreScript(value=SCRIPT)
+    GEMTransform().apply(script, base_job)
+
+    assert script.value == create_expected_script("12345-12347")
+
+
+def test_gem_transform_non_contiguous_runs(base_job, create_expected_script):
+    """Test GEMTransform with non-contiguous runs."""
+    base_job.inputs["runno"] = [12345, 12347, 12349]
+    script = PreScript(value=SCRIPT)
+    GEMTransform().apply(script, base_job)
+
+    assert script.value == create_expected_script("12345,12347,12349")
+
+
+def test_gem_transform_list_length_one(base_job, create_expected_script):
+    """Test GEMTransform with list containing a single run."""
+    base_job.inputs["runno"] = [12345]
+    script = PreScript(value=SCRIPT)
+    GEMTransform().apply(script, base_job)
+
+    assert script.value == create_expected_script("12345")
+
+
+def test_gem_transform_apply(base_job):
+    """Test GEMTransform only modifies expected lines and leaves others unchanged."""
+    transform = GEMTransform()
+    script = PreScript(value=SCRIPT)
+    original_lines = script.value.splitlines()
+
+    transform.apply(script, base_job)
+
+    updated_lines = script.value.splitlines()
+    assert len(original_lines) == len(updated_lines)
+
+    for index, line in enumerate(updated_lines):
+        if line.startswith("mode ="):
+            assert line == 'mode = "transmission"'
+        elif line.startswith("input_mode ="):
+            assert line == 'input_mode = "raw"'
+        elif line.startswith("vanadium_runno ="):
+            assert line == "vanadium_runno = 12345"
+        elif line.startswith("runno ="):
+            assert line == "runno = 12345"
+        elif line.startswith("calibration_dir ="):
+            assert line == "calibration_dir = /path/to/cal"
+        elif line.startswith("splined_vanadium_dir ="):
+            assert line == 'splined_vanadium_dir = "/path/to/splined"'
+        elif line.startswith("config_file ="):
+            assert line == 'config_file = "/path/to/config"'
+        elif line.startswith("output_dir = "):
+            assert line == 'output_dir = "/path/to/output"'
+        else:
+            assert line == original_lines[index]