feat: Add ability to use datetime objects or time-shift strings in TimestampRange (#2432)

Author: MortGron

cognite/client/_api/simulators/models_revisions.py

@@ -73,8 +73,8 @@ async def list(
                 >>> res = client.simulators.models.revisions.list(
                 ...     model_external_ids=["model1", "model2"],
                 ...     all_versions=True,
-                ...     created_time=TimestampRange(min=0, max=1000000),
-                ...     last_updated_time=TimestampRange(min=0, max=1000000),
+                ...     created_time=TimestampRange(min="1d-ago", max="now"),
+                ...     last_updated_time=TimestampRange(min="1d-ago", max="now"),
                 ...     sort=PropertySort(order="asc", property="createdTime"),
                 ...     limit=10,
                 ... )

cognite/client/_api/simulators/runs.py

@@ -197,8 +197,8 @@ async def list(
             Filter runs by time ranges:
                 >>> from cognite.client.data_classes.shared import TimestampRange
                 >>> res = client.simulators.runs.list(
-                ...     created_time=TimestampRange(min=0, max=1_700_000_000_000),
-                ...     simulation_time=TimestampRange(min=0, max=1_700_000_000_000),
+                ...     created_time=TimestampRange(min="1d-ago", max="now"),
+                ...     simulation_time=TimestampRange(min="1d-ago", max="now"),
                 ... )
         """
 

cognite/client/_sync_api/simulators/models_revisions.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-143e2ea035bde172593740ba9733e0e2
+03b4ddd2ad52415dca0fe62f12b4a3c1
 This file is auto-generated from the Async API modules, - do not edit manually!
 ===============================================================================
 """
@@ -72,8 +72,8 @@ def list(
                 >>> res = client.simulators.models.revisions.list(
                 ...     model_external_ids=["model1", "model2"],
                 ...     all_versions=True,
-                ...     created_time=TimestampRange(min=0, max=1000000),
-                ...     last_updated_time=TimestampRange(min=0, max=1000000),
+                ...     created_time=TimestampRange(min="1d-ago", max="now"),
+                ...     last_updated_time=TimestampRange(min="1d-ago", max="now"),
                 ...     sort=PropertySort(order="asc", property="createdTime"),
                 ...     limit=10,
                 ... )

cognite/client/_sync_api/simulators/runs.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-183a4d1e3ee0ff2c01d10dadfbdda9a0
+c2ccaaf731fd07a67787a512a5a9acd2
 This file is auto-generated from the Async API modules, - do not edit manually!
 ===============================================================================
 """
@@ -185,8 +185,8 @@ def list(
             Filter runs by time ranges:
                 >>> from cognite.client.data_classes.shared import TimestampRange
                 >>> res = client.simulators.runs.list(
-                ...     created_time=TimestampRange(min=0, max=1_700_000_000_000),
-                ...     simulation_time=TimestampRange(min=0, max=1_700_000_000_000),
+                ...     created_time=TimestampRange(min="1d-ago", max="now"),
+                ...     simulation_time=TimestampRange(min="1d-ago", max="now"),
                 ... )
         """
         return run_sync(

cognite/client/data_classes/shared.py

@@ -1,25 +1,29 @@
 from __future__ import annotations
 
 from collections.abc import Collection, Sequence
+from datetime import datetime
 from typing import Any, Literal
 
 from typing_extensions import Self
 
 from cognite.client.data_classes._base import CogniteFilter, CogniteResource, UnknownCogniteResource
+from cognite.client.utils._time import timestamp_to_ms
 
 
 class TimestampRange(CogniteResource):
     """Range between two timestamps.
 
     Args:
-        max (int | None): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
-        min (int | None): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
+        max (int | float | str | datetime | None): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds, a string in time-shift format or a datetime object.
+        min (int | float | str | datetime | None): The number of milliseconds since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds, a string in time-shift format or a datetime object.
         **_ (Any): No description.
     """
 
-    def __init__(self, max: int | None = None, min: int | None = None, **_: Any) -> None:
-        self.max = max
-        self.min = min
+    def __init__(
+        self, max: int | float | str | datetime | None = None, min: int | float | str | datetime | None = None, **_: Any
+    ) -> None:
+        self.max = timestamp_to_ms(max) if max is not None else None
+        self.min = timestamp_to_ms(min) if min is not None else None
 
     @classmethod
     def _load(cls, resource: dict[str, Any]) -> Self:

tests/tests_unit/test_api/test_data_sets.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
 
 import pytest
@@ -65,8 +66,11 @@ def test_retrieve_multiple(
         assert isinstance(res, DataSetList)
         assert [example_data_set] == res.dump(camel_case=True)
 
-    def test_list_with_timestamp_range(self, cognite_client: CogniteClient, mock_ds_response: HTTPXMock) -> None:
-        cognite_client.data_sets.list(created_time=TimestampRange(min=20))
+    @pytest.mark.parametrize("min_time", [20, datetime.fromtimestamp(20 / 1000, timezone.utc)])
+    def test_list_with_timestamp_range(
+        self, cognite_client: CogniteClient, mock_ds_response: HTTPXMock, min_time: int | datetime
+    ) -> None:
+        cognite_client.data_sets.list(created_time=TimestampRange(min=min_time))
         assert 20 == jsgz_load(mock_ds_response.get_requests()[0].content)["filter"]["createdTime"]["min"]
         assert "max" not in jsgz_load(mock_ds_response.get_requests()[0].content)["filter"]["createdTime"]
 

tests/tests_unit/test_data_classes/test_timestamp_range.py

@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+from datetime import datetime, timedelta, timezone
+
+import pytest
+
 from cognite.client.data_classes import AggregateResultItem, TimestampRange
 
 
@@ -27,3 +31,19 @@ def test_camels(self) -> None:
         ag = AggregateResultItem(child_count=23, depth=1, path=[])
         assert 23 == ag.child_count
         assert {"childCount": 23, "depth": 1, "path": []} == ag.dump(camel_case=True)
+
+    def test_datetime(self) -> None:
+        min_time = datetime.fromtimestamp(1767222000, timezone.utc)  # 2026-01-01 00:00:00 GMT+01
+        max_time = datetime.fromtimestamp(1767308400, timezone.utc)  # 2026-01-02 00:00:00 GMT+01
+        tsr = TimestampRange(min=min_time, max=max_time)
+        assert tsr.min == 1767222000000
+        assert tsr.max == 1767308400000
+        assert {"min": 1767222000000, "max": 1767308400000} == tsr.dump()
+
+    def test_time_shift_string(self) -> None:
+        now_time = int(datetime.now(timezone.utc).timestamp() * 1000)
+        day_ago_time = now_time - timedelta(days=1) // timedelta(milliseconds=1)
+        tsr = TimestampRange(min="1d-ago", max="now")
+        # NOTE: Using approx because of the small time difference between calls
+        assert tsr.min == pytest.approx(day_ago_time, 100)
+        assert tsr.max == pytest.approx(now_time, 100)