feat: add helper methods to access simulator quantities and units (#2305)

Author: evertoncolling

CHANGELOG.md

@@ -24,6 +24,7 @@ As of 2025-08-29, changes are grouped as follows
 ## [Unreleased]
 ### Added
 - [alpha] Support for `created_time` and `simulation_time` filters in `client.simulators.runs.list()` to filter simulation runs by timestamp ranges.
+- Added `get_quantities()` and `get_units()` methods to the `Simulator` class for easier access to simulator unit quantities and their units.
 
 ## [7.82.1] - 2025-08-28
 ### Fixed

cognite/client/data_classes/simulators/simulators.py

@@ -108,6 +108,38 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
 
         return output
 
+    def get_quantities(self) -> list[str]:
+        """Get a list of quantity names available for this simulator.
+
+        Returns:
+            list[str]: List of quantity names from the simulator's unit_quantities.
+        """
+        if not self.unit_quantities:
+            return []
+        return [quantity.name for quantity in self.unit_quantities]
+
+    def get_units(self, quantity: str) -> list[str]:
+        """Get a list of unit names for a specific quantity.
+
+        Args:
+            quantity (str): The name of the quantity to get units for.
+
+        Returns:
+            list[str]: List of unit names for the specified quantity.
+
+        Raises:
+            ValueError: If the specified quantity does not exist for this simulator.
+        """
+        if not self.unit_quantities:
+            raise ValueError(f"Quantity '{quantity}' not found. This simulator has no unit quantities defined.")
+
+        for q in self.unit_quantities:
+            if q.name == quantity:
+                return [unit.name for unit in q.units]
+
+        available_quantities = self.get_quantities()
+        raise ValueError(f"Quantity '{quantity}' not found. Available quantities: {', '.join(available_quantities)}")
+
 
 @dataclass
 class SimulatorUnitEntry(CogniteObject):
@@ -149,7 +181,9 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
 
     @classmethod
     def _load_list(
-        cls, resource: dict[str, Any] | list[dict[str, Any]], cognite_client: CogniteClient | None = None
+        cls,
+        resource: dict[str, Any] | list[dict[str, Any]],
+        cognite_client: CogniteClient | None = None,
     ) -> list[SimulatorModelType]:
         if isinstance(resource, dict):
             return [cls._load(resource, cognite_client)]
@@ -175,7 +209,9 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
 
     @classmethod
     def _load_list(
-        cls, resource: dict[str, Any] | list[dict[str, Any]], cognite_client: CogniteClient | None = None
+        cls,
+        resource: dict[str, Any] | list[dict[str, Any]],
+        cognite_client: CogniteClient | None = None,
     ) -> list[SimulatorQuantity]:
         if isinstance(resource, dict):
             return [cls._load(resource, cognite_client)]
@@ -261,7 +297,9 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
 
     @classmethod
     def _load_list(
-        cls, resource: dict[str, Any] | list[dict[str, Any]], cognite_client: CogniteClient | None = None
+        cls,
+        resource: dict[str, Any] | list[dict[str, Any]],
+        cognite_client: CogniteClient | None = None,
     ) -> list[SimulatorModelDependency]:
         if isinstance(resource, dict):
             return [cls._load(resource, cognite_client)]
@@ -291,7 +329,9 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
 
     @classmethod
     def _load_list(
-        cls, resource: dict[str, Any] | list[dict[str, Any]], cognite_client: CogniteClient | None = None
+        cls,
+        resource: dict[str, Any] | list[dict[str, Any]],
+        cognite_client: CogniteClient | None = None,
     ) -> list[SimulatorStep]:
         if isinstance(resource, dict):
             return [cls._load(resource, cognite_client)]

tests/tests_unit/test_data_classes/test_simulators.py

@@ -4,6 +4,11 @@
     SimulatorModelDependencyFileId,
     SimulatorModelRevisionDependency,
 )
+from cognite.client.data_classes.simulators.simulators import (
+    Simulator,
+    SimulatorQuantity,
+    SimulatorUnitEntry,
+)
 
 
 class TestSimulatorModelRevisionDependency:
@@ -28,3 +33,102 @@ def test_load_list(self, input_data):
         assert result[0].arguments == {"fieldA": "valueA"}
         assert isinstance(result[0].file, SimulatorModelDependencyFileId)
         assert result[1].file.id == 2222
+
+
+class TestSimulator:
+    @pytest.fixture
+    def simulator_with_quantities(self):
+        """Create a simulator with unit quantities for testing."""
+        unit_m = SimulatorUnitEntry(label="Meters", name="m")
+        unit_ft = SimulatorUnitEntry(label="Feet", name="ft")
+        unit_kg = SimulatorUnitEntry(label="Kilograms", name="kg")
+        unit_lb = SimulatorUnitEntry(label="Pounds", name="lb")
+        unit_c = SimulatorUnitEntry(label="Celsius", name="C")
+        unit_f = SimulatorUnitEntry(label="Fahrenheit", name="F")
+
+        quantity_length = SimulatorQuantity(name="length", label="Length", units=[unit_m, unit_ft])
+        quantity_mass = SimulatorQuantity(name="mass", label="Mass", units=[unit_kg, unit_lb])
+        quantity_temp = SimulatorQuantity(name="temperature", label="Temperature", units=[unit_c, unit_f])
+
+        return Simulator(
+            external_id="test-simulator",
+            id=1,
+            name="Test Simulator",
+            file_extension_types=[".sim"],
+            unit_quantities=[quantity_length, quantity_mass, quantity_temp],
+        )
+
+    @pytest.fixture
+    def simulator_without_quantities(self):
+        """Create a simulator without unit quantities for testing."""
+        return Simulator(
+            external_id="empty-simulator",
+            id=2,
+            name="Empty Simulator",
+            file_extension_types=[".sim"],
+            unit_quantities=None,
+        )
+
+    def test_get_quantities_with_data(self, simulator_with_quantities):
+        """Test get_quantities returns correct list when quantities exist."""
+        quantities = simulator_with_quantities.get_quantities()
+        assert quantities == ["length", "mass", "temperature"]
+        assert len(quantities) == 3
+        assert all(isinstance(q, str) for q in quantities)
+
+    def test_get_quantities_empty(self, simulator_without_quantities):
+        """Test get_quantities returns empty list when no quantities exist."""
+        quantities = simulator_without_quantities.get_quantities()
+        assert quantities == []
+        assert len(quantities) == 0
+
+    def test_get_units_valid_quantity(self, simulator_with_quantities):
+        """Test get_units returns correct units for valid quantities."""
+        # Test length units
+        length_units = simulator_with_quantities.get_units("length")
+        assert length_units == ["m", "ft"]
+        assert len(length_units) == 2
+
+        # Test mass units
+        mass_units = simulator_with_quantities.get_units("mass")
+        assert mass_units == ["kg", "lb"]
+        assert len(mass_units) == 2
+
+        # Test temperature units
+        temp_units = simulator_with_quantities.get_units("temperature")
+        assert temp_units == ["C", "F"]
+        assert len(temp_units) == 2
+
+    def test_get_units_invalid_quantity(self, simulator_with_quantities):
+        """Test get_units raises ValueError for non-existent quantity."""
+        with pytest.raises(ValueError) as exc_info:
+            simulator_with_quantities.get_units("pressure")
+
+        error_msg = str(exc_info.value)
+        assert "Quantity 'pressure' not found" in error_msg
+        assert "Available quantities: length, mass, temperature" in error_msg
+
+    def test_get_units_no_quantities(self, simulator_without_quantities):
+        """Test get_units raises ValueError when simulator has no quantities."""
+        with pytest.raises(ValueError) as exc_info:
+            simulator_without_quantities.get_units("any_quantity")
+
+        error_msg = str(exc_info.value)
+        assert "Quantity 'any_quantity' not found" in error_msg
+        assert "This simulator has no unit quantities defined" in error_msg
+
+    def test_get_units_empty_units_list(self):
+        """Test get_units with quantity that has empty units list."""
+        empty_quantity = SimulatorQuantity(name="empty", label="Empty Quantity", units=[])
+
+        simulator = Simulator(
+            external_id="test-sim",
+            id=3,
+            name="Test",
+            file_extension_types=[".sim"],
+            unit_quantities=[empty_quantity],
+        )
+
+        units = simulator.get_units("empty")
+        assert units == []
+        assert len(units) == 0