Fix retrieve_dataframe_in_tz raising on certain datetime inputs (#…

…1550)

CHANGELOG.md

@@ -17,6 +17,11 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [7.5.7] - 2023-12-12
+### Fixed
+- Certain combinations of `start`/`end` and `granularity` would cause `retrieve_dataframe_in_tz` to raise due to
+  a bug in the calender-arithmetic (`MonthAligner`).
+
 ## [7.5.6] - 2023-12-11
 ### Added
 - Missing legacy scopes for `Capability`: `LegacySpaceScope` and `LegacyDataModelScope`.

cognite/client/_version.py

@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.5.6"
+__version__ = "7.5.7"
 __api_subversion__ = "V20220125"

cognite/client/utils/_time.py

@@ -320,8 +320,8 @@ def ceil(cls, date: datetime) -> datetime:
         """
         if date == datetime(year=date.year, month=date.month, day=1, tzinfo=date.tzinfo):
             return date
-        extra, month = divmod(date.month + 1, 12)
-        return cls.normalize(date.replace(year=date.year + extra, month=month, day=1))
+        extra, month = divmod(date.month, 12)  # this works because month is one-indexed
+        return cls.normalize(date.replace(year=date.year + extra, month=month + 1, day=1))
 
     @classmethod
     def floor(cls, date: datetime) -> datetime:
@@ -333,8 +333,12 @@ def units_between(cls, start: datetime, end: datetime) -> int:
 
     @classmethod
     def add_units(cls, date: datetime, units: int) -> datetime:
-        extra_years, month = divmod(date.month + units, 12)
-        return date.replace(year=date.year + extra_years, month=month)
+        """
+        Adds 'units' number of months to 'date', ignoring timezone. The resulting date might not be valid,
+        for example Jan 29 + 1 unit in a non-leap year. In such cases, datetime will raise a ValueError.
+        """
+        extra_years, month = divmod(date.month + units - 1, 12)
+        return date.replace(year=date.year + extra_years, month=month + 1)
 
 
 class QuarterAligner(DateTimeAligner):
@@ -376,6 +380,10 @@ def units_between(cls, start: datetime, end: datetime) -> int:
 
     @classmethod
     def add_units(cls, date: datetime, units: int) -> datetime:
+        """
+        Adds 'units' number of quarters to 'date', ignoring timezone. The resulting date might not be valid,
+        for example Jan 31 + 1 unit, as April only has 30 days. In such cases, datetime will raise a ValueError.
+        """
         extra_years, month = divmod(date.month + 3 * units, 12)
         return date.replace(year=date.year + extra_years, month=month)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "7.5.6"
+version = "7.5.7"
 description = "Cognite Python SDK"
 readme = "README.md"
 documentation = "https://cognite-sdk-python.readthedocs-hosted.com"

tests/tests_unit/test_utils/test_time.py

@@ -14,6 +14,7 @@
 from cognite.client.utils._time import (
     MAX_TIMESTAMP_MS,
     MIN_TIMESTAMP_MS,
+    MonthAligner,
     align_large_granularity,
     align_start_and_end_for_granularity,
     convert_and_isoformat_time_attrs,
@@ -664,3 +665,48 @@ def test_pandas_date_range_tz_ambiguous_time_error():
 
         # Assert
         assert len(index) == expected_length
+
+
+class TestDateTimeAligner:
+    # TODO: DayAligner
+    # TODO: WeekAligner
+    # TODO: MonthAligner
+    # TODO: QuarterAligner
+    # TODO: YearAligner
+
+    @pytest.mark.parametrize(
+        "dt, expected",
+        (
+            (datetime(2023, 11, 1), datetime(2023, 11, 1)),
+            (datetime(2023, 10, 15), datetime(2023, 11, 1)),
+            (datetime(2023, 12, 15), datetime(2024, 1, 1)),
+            (datetime(2024, 1, 10), datetime(2024, 2, 1)),
+            # Bug prior to 7.5.7 would cause this to raise:
+            (datetime(2023, 11, 2), datetime(2023, 12, 1)),
+        ),
+    )
+    def test_month_aligner__ceil(self, dt, expected):
+        assert expected == MonthAligner.ceil(dt)
+
+    def test_month_aligner_ceil__invalid_date(self):
+        with pytest.raises(ValueError, match="^day is out of range for month$"):
+            MonthAligner.add_units(datetime(2023, 7, 31), 2)  # sept has 30 days
+
+    @pytest.mark.parametrize(
+        "dt, n_units, expected",
+        (
+            (datetime(2023, 7, 2), 12, datetime(2024, 7, 2)),
+            (datetime(2023, 7, 2), 12 * 12, datetime(2035, 7, 2)),
+            (datetime(2023, 7, 2), -12 * 2, datetime(2021, 7, 2)),
+            # Bug prior to 7.5.7 would cause these to raise:
+            (datetime(2023, 11, 15), 1, datetime(2023, 12, 15)),
+            (datetime(2023, 12, 15), 0, datetime(2023, 12, 15)),
+            (datetime(2024, 1, 15), -1, datetime(2023, 12, 15)),
+        ),
+    )
+    def test_month_aligner__add_unites(self, dt, n_units, expected):
+        assert expected == MonthAligner.add_units(dt, n_units)
+
+    def test_month_aligner_add_unites__invalid_date(self):
+        with pytest.raises(ValueError, match="^day is out of range for month$"):
+            MonthAligner.add_units(datetime(2023, 1, 29), 1)  # 2023 = non-leap year