feat: add sort parameter to simulators API list and iteration methods (#2355)

Author: evertoncolling

cognite/client/_api/simulators/models.py

@@ -189,6 +189,7 @@ def __call__(
             resource_cls=SimulatorModel,
             method="POST",
             filter=model_filter.dump(),
+            sort=[PropertySort.load(sort).dump()] if sort else None,
             chunk_size=chunk_size,
             limit=limit,
         )

cognite/client/_api/simulators/models_revisions.py

@@ -220,6 +220,7 @@ def __call__(
             resource_cls=SimulatorModelRevision,
             method="POST",
             filter=model_revisions_filter.dump(),
+            sort=[PropertySort.load(sort).dump()] if sort else None,
             chunk_size=chunk_size,
             limit=limit,
         )

cognite/client/_api/simulators/routines.py

@@ -54,6 +54,7 @@ def __call__(
         chunk_size: int,
         model_external_ids: Sequence[str] | None = None,
         simulator_integration_external_ids: Sequence[str] | None = None,
+        sort: PropertySort | None = None,
         limit: int | None = None,
     ) -> Iterator[SimulatorRoutineList]: ...
 
@@ -63,6 +64,7 @@ def __call__(
         chunk_size: None = None,
         model_external_ids: Sequence[str] | None = None,
         simulator_integration_external_ids: Sequence[str] | None = None,
+        sort: PropertySort | None = None,
         limit: int | None = None,
     ) -> Iterator[SimulatorRoutine]: ...
 
@@ -71,6 +73,7 @@ def __call__(
         chunk_size: int | None = None,
         model_external_ids: Sequence[str] | None = None,
         simulator_integration_external_ids: Sequence[str] | None = None,
+        sort: PropertySort | None = None,
         limit: int | None = None,
     ) -> Iterator[SimulatorRoutine] | Iterator[SimulatorRoutineList]:
         """Iterate over simulator routines
@@ -81,6 +84,7 @@ def __call__(
             chunk_size (int | None): Number of simulator routines to return in each chunk. Defaults to yielding one simulator routine a time.
             model_external_ids (Sequence[str] | None): Filter on model external ids.
             simulator_integration_external_ids (Sequence[str] | None): Filter on simulator integration external ids.
+            sort (PropertySort | None): The criteria to sort by.
             limit (int | None): Maximum number of simulator routines to return. Defaults to return all items.
 
         Returns:
@@ -96,6 +100,7 @@ def __call__(
             resource_cls=SimulatorRoutine,
             method="POST",
             filter=routines_filter.dump(),
+            sort=[PropertySort.load(sort).dump()] if sort else None,
             chunk_size=chunk_size,
             limit=limit,
         )

cognite/client/_api/simulators/runs.py

@@ -6,7 +6,7 @@
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
 from cognite.client.data_classes.shared import TimestampRange
-from cognite.client.data_classes.simulators.filters import SimulatorRunsFilter
+from cognite.client.data_classes.simulators.filters import SimulationRunsSort, SimulatorRunsFilter
 from cognite.client.data_classes.simulators.runs import (
     SimulationRun,
     SimulationRunDataList,
@@ -67,6 +67,7 @@ def __call__(
         model_revision_external_ids: SequenceNotStr[str] | None = None,
         created_time: TimestampRange | None = None,
         simulation_time: TimestampRange | None = None,
+        sort: SimulationRunsSort | None = None,
     ) -> Iterator[SimulationRunList]: ...
 
     @overload
@@ -84,6 +85,7 @@ def __call__(
         model_revision_external_ids: SequenceNotStr[str] | None = None,
         created_time: TimestampRange | None = None,
         simulation_time: TimestampRange | None = None,
+        sort: SimulationRunsSort | None = None,
     ) -> Iterator[SimulationRun]: ...
 
     def __call__(
@@ -100,6 +102,7 @@ def __call__(
         model_revision_external_ids: SequenceNotStr[str] | None = None,
         created_time: TimestampRange | None = None,
         simulation_time: TimestampRange | None = None,
+        sort: SimulationRunsSort | None = None,
     ) -> Iterator[SimulationRun] | Iterator[SimulationRunList]:
         """Iterate over simulation runs
 
@@ -118,6 +121,7 @@ def __call__(
             model_revision_external_ids (SequenceNotStr[str] | None): Filter by model revision external ids
             created_time (TimestampRange | None): Filter by created time
             simulation_time (TimestampRange | None): Filter by simulation time
+            sort (SimulationRunsSort | None): The criteria to sort by.
 
         Returns:
             Iterator[SimulationRun] | Iterator[SimulationRunList]: yields Simulation Run one by one if chunk is not specified, else SimulatorRunsList objects.
@@ -141,6 +145,7 @@ def __call__(
             resource_cls=SimulationRun,
             method="POST",
             filter=filter_runs.dump(),
+            sort=[SimulationRunsSort.load(sort).dump()] if sort else None,
             chunk_size=chunk_size,
             limit=limit,
         )
@@ -158,6 +163,7 @@ def list(
         model_revision_external_ids: SequenceNotStr[str] | None = None,
         created_time: TimestampRange | None = None,
         simulation_time: TimestampRange | None = None,
+        sort: SimulationRunsSort | None = None,
     ) -> SimulationRunList:
         """`Filter simulation runs <https://developer.cognite.com/api#tag/Simulation-Runs/operation/filter_simulation_runs_simulators_runs_list_post>`_
 
@@ -175,6 +181,7 @@ def list(
             model_revision_external_ids (SequenceNotStr[str] | None): Filter by model revision external ids
             created_time (TimestampRange | None): Filter by created time
             simulation_time (TimestampRange | None): Filter by simulation time
+            sort (SimulationRunsSort | None): The criteria to sort by.
 
         Returns:
             SimulationRunList: List of simulation runs
@@ -218,6 +225,7 @@ def list(
             resource_cls=SimulationRun,
             list_cls=SimulationRunList,
             filter=filter_runs.dump(),
+            sort=[SimulationRunsSort.load(sort).dump()] if sort else None,
         )
 
     @overload

cognite/client/data_classes/simulators/__init__.py

@@ -2,6 +2,7 @@
 
 from cognite.client.data_classes.simulators.filters import (
     PropertySort,
+    SimulationRunsSort,
     SimulatorIntegrationFilter,
     SimulatorModelRevisionsFilter,
     SimulatorModelsFilter,
@@ -37,12 +38,12 @@
 )
 
 __all__ = [
-    "PropertySort",
     "PropertySort",
     "SimulationRun",
     "SimulationRunList",
     "SimulationRunWrite",
     "SimulationRunWriteList",
+    "SimulationRunsSort",
     "Simulator",
     "SimulatorIntegration",
     "SimulatorIntegrationFilter",

cognite/client/data_classes/simulators/filters.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from cognite.client.data_classes._base import CogniteFilter, CogniteSort
 from cognite.client.data_classes.shared import TimestampRange
+from cognite.client.utils._text import to_camel_case
 from cognite.client.utils.useful_types import SequenceNotStr
 
 
@@ -91,14 +92,30 @@ def __init__(
 class PropertySort(CogniteSort):
     def __init__(
         self,
-        property: Literal["createdTime"] = "createdTime",
+        property: Literal["created_time"] = "created_time",
         order: Literal["asc", "desc"] = "asc",
     ):
         super().__init__(property, order)
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         dumped = super().dump(camel_case=camel_case)
-        dumped["property"] = self.property
+        prop = cast(str, self.property)
+        dumped["property"] = to_camel_case(prop) if camel_case else prop
+        return dumped
+
+
+class SimulationRunsSort(CogniteSort):
+    def __init__(
+        self,
+        property: Literal["created_time", "simulation_time"] = "created_time",
+        order: Literal["asc", "desc"] = "asc",
+    ):
+        super().__init__(property, order)
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        dumped = super().dump(camel_case=camel_case)
+        prop = cast(str, self.property)
+        dumped["property"] = to_camel_case(prop) if camel_case else prop
         return dumped
 
 

tests/tests_integration/test_api/test_simulators/test_models.py

@@ -41,6 +41,7 @@ def test_list_models(self, cognite_client: CogniteClient, seed_resource_names: R
         for model in cognite_client.simulators.models(
             limit=2,
             simulator_external_ids=[seed_resource_names.simulator_external_id],
+            sort=PropertySort(order="asc", property="created_time"),
         ):
             assert model.created_time is not None
             model_ids.append(model.id)
@@ -72,7 +73,10 @@ def test_list_model_revisions(self, cognite_client: CogniteClient, seed_resource
         )
 
         model_revision_ids = []
-        for revision in cognite_client.simulators.models.revisions(limit=2):
+        for revision in cognite_client.simulators.models.revisions(
+            limit=2,
+            sort=PropertySort(order="desc", property="created_time"),
+        ):
             assert revision.created_time is not None
             model_revision_ids.append(revision.id)
 
@@ -105,13 +109,13 @@ def test_list_model_revisions_filtering_sort(
         self, cognite_client: CogniteClient, seed_resource_names: ResourceNames
     ) -> None:
         revisions_asc = cognite_client.simulators.models.revisions.list(
-            sort=PropertySort(order="asc", property="createdTime"),
+            sort=PropertySort(order="asc", property="created_time"),
             model_external_ids=[seed_resource_names.simulator_model_external_id],
             all_versions=True,
         )
 
         revisions_desc = cognite_client.simulators.models.revisions.list(
-            sort=PropertySort(order="desc", property="createdTime"),
+            sort=PropertySort(order="desc", property="created_time"),
             model_external_ids=[seed_resource_names.simulator_model_external_id],
             all_versions=True,
         )

tests/tests_integration/test_api/test_simulators/test_routine_revisions.py

@@ -38,7 +38,7 @@ def test_list_and_filtering_routine_revisions(
         revisions_by_model: list[SimulatorRoutineRevision] = []
 
         for revision in cognite_client.simulators.routines.revisions(
-            sort=PropertySort(order="asc", property="createdTime"),
+            sort=PropertySort(order="asc", property="created_time"),
             model_external_ids=[model_external_id],
             all_versions=True,
             include_all_fields=True,

tests/tests_integration/test_api/test_simulators/test_routines.py

@@ -1,3 +1,5 @@
+from itertools import pairwise
+
 import pytest
 
 from cognite.client._cognite_client import CogniteClient
@@ -46,9 +48,22 @@ def test_sort(
 
         routines_asc = cognite_client.simulators.routines.list(
             simulator_integration_external_ids=[simulator_integration_unique_external_id],
-            sort=PropertySort(order="asc", property="createdTime"),
+            sort=PropertySort(order="asc", property="created_time"),
         )
 
         assert len(routines_asc) > 1
         for i in range(1, len(routines_asc)):
             assert routines_asc[i].created_time >= routines_asc[i - 1].created_time
+
+        # Test iterator with sort
+        routines_iter = list(
+            cognite_client.simulators.routines(
+                simulator_integration_external_ids=[simulator_integration_unique_external_id],
+                sort=PropertySort(order="asc", property="created_time"),
+                limit=3,
+            )
+        )
+
+        assert len(routines_iter) == 3
+        for prev, curr in pairwise(routines_iter):
+            assert curr.created_time >= prev.created_time

tests/tests_integration/test_api/test_simulators/test_runs.py

@@ -1,10 +1,12 @@
 import asyncio
 import time
+from itertools import pairwise
 
 import pytest
 
 from cognite.client._cognite_client import CogniteClient
 from cognite.client.data_classes import TimestampRange
+from cognite.client.data_classes.simulators.filters import SimulationRunsSort
 from cognite.client.data_classes.simulators.routine_revisions import SimulatorRoutineRevision
 from cognite.client.data_classes.simulators.runs import (
     SimulationInput,
@@ -181,6 +183,33 @@ def test_list_filtering_timestamp_ranges(
         )
         assert len(runs_empty_range) == 0
 
+    def test_list_sort(self, cognite_client: CogniteClient, seed_resource_names: ResourceNames) -> None:
+        routine_external_id = seed_resource_names.simulator_routine_external_id
+
+        # Test list with sort by createdTime
+        runs_asc = cognite_client.simulators.runs.list(
+            routine_external_ids=[routine_external_id],
+            sort=SimulationRunsSort(order="asc", property="created_time"),
+            limit=5,
+        )
+
+        if len(runs_asc) > 1:
+            for prev, curr in pairwise(runs_asc):
+                assert curr.created_time >= prev.created_time
+
+        # Test iterator with sort by createdTime descending
+        runs_iter = list(
+            cognite_client.simulators.runs(
+                routine_external_ids=[routine_external_id],
+                sort=SimulationRunsSort(order="desc", property="created_time"),
+                limit=3,
+            )
+        )
+
+        if len(runs_iter) > 1:
+            for prev, curr in pairwise(runs_iter):
+                assert curr.created_time <= prev.created_time
+
     def test_list_run_data(self, cognite_client: CogniteClient, seed_resource_names: ResourceNames) -> None:
         routine_external_id = seed_resource_names.simulator_routine_external_id
         created_run = cognite_client.simulators.runs.create(