apply black

Signed-off-by: Victor Garcia Reolid <[email protected]>
SeitaBV · Mar 20, 2024 · 79bac48 · 79bac48
1 parent 5b30af0
commit 79bac48
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 38 deletions.
diff --git a/timely_beliefs/sensors/func_store/knowledge_horizons.py b/timely_beliefs/sensors/func_store/knowledge_horizons.py
@@ -7,7 +7,10 @@
 
 import pandas as pd
 
-from timely_beliefs.sensors.func_store.utils import datetime_x_days_ago_at_y_oclock, x_years_ago_at_date_datetime
+from timely_beliefs.sensors.func_store.utils import (
+    datetime_x_days_ago_at_y_oclock,
+    x_years_ago_at_date_datetime,
+)
 
 
 def at_date(
@@ -31,18 +34,17 @@ def at_date(
     return event_start - knowledge_time.astimezone(event_start.tzinfo)
 
 
-
 def x_years_ago_at_date(
     event_start: datetime | pd.DatetimeIndex,
     day: int,
-    month : int,
-    x : int = 1,
+    month: int,
+    x: int = 1,
     get_bounds: bool = False,
     z: str | None = None,
 ) -> timedelta | pd.TimedeltaIndex | tuple[timedelta, timedelta]:
     """Compute the sensor's knowledge horizon to represent the event could be known since some date, `x` years ago.
-    
-    For example, it can be used for a tax rate that changes annually and with a known publication date. 
+
+    For example, it can be used for a tax rate that changes annually and with a known publication date.
 
     :param event_start:     Start of the event, used as an anchor for determining the knowledge horizon.
     :param day:             Reference day of the month of the annual date to compare against.
@@ -54,11 +56,16 @@ def x_years_ago_at_date(
     """
     if get_bounds:
         return timedelta.min, timedelta.days(366 * x)
-        
+
     if isinstance(event_start, datetime):
         return x_years_ago_at_date_datetime(event_start, day, month, x, z)
     else:
-        return event_start.map(lambda _event_start: x_years_ago_at_date_datetime(_event_start,day,month, x, z))
+        return event_start.map(
+            lambda _event_start: x_years_ago_at_date_datetime(
+                _event_start, day, month, x, z
+            )
+        )
+
 
 def ex_post(
     event_resolution: timedelta,

diff --git a/timely_beliefs/sensors/func_store/test_knowledge_horizons.py b/timely_beliefs/sensors/func_store/test_knowledge_horizons.py
@@ -58,67 +58,89 @@ def test_fixed_knowledge_time():
 def test_x_years_ago_at_date():
     """Check definition of knowledge horizon for events known at a fixed date annually."""
 
-    knowledge_func_params = dict(month=11, day=20)    
-    
+    knowledge_func_params = dict(month=11, day=20)
+
     # Events that occur before the reference
     # year 2024 is leap
     assert x_years_ago_at_date(
-        event_start=datetime(2024, 11, 19, 1, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=365, hours=1) # 366 days - 1
+        event_start=datetime(2024, 11, 19, 1, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=365, hours=1
+    )  # 366 days - 1
 
     # year 2025 is not leap, but 2024 is
     assert x_years_ago_at_date(
-        event_start=datetime(2025, 11, 19, 2, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=364, hours=2) # 365 - 1
+        event_start=datetime(2025, 11, 19, 2, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=364, hours=2
+    )  # 365 - 1
 
     # year 2023 is not leap and 2022 neither
     assert x_years_ago_at_date(
-        event_start=datetime(2022, 11, 19, 2, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=364, hours=2) # 365 - 1
+        event_start=datetime(2022, 11, 19, 2, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=364, hours=2
+    )  # 365 - 1
 
     # Events that occur after the reference
     assert x_years_ago_at_date(
-        event_start=datetime(2021, 11, 21, 3, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=366, hours=3) # 365 + 1
+        event_start=datetime(2021, 11, 21, 3, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=366, hours=3
+    )  # 365 + 1
 
     assert x_years_ago_at_date(
-        event_start=datetime(2021, 11, 21, 4, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=366, hours=4) # 365 + 1
+        event_start=datetime(2021, 11, 21, 4, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=366, hours=4
+    )  # 365 + 1
 
     assert x_years_ago_at_date(
-        event_start=datetime(2020, 11, 21, 4, tzinfo=utc),
-        **knowledge_func_params
-    ) == timedelta(days=367, hours=4) # 366 (leap year) + 1
+        event_start=datetime(2020, 11, 21, 4, tzinfo=utc), **knowledge_func_params
+    ) == timedelta(
+        days=367, hours=4
+    )  # 366 (leap year) + 1
 
     # Test a Daylight Savings Transition
-    knowledge_func_params_dst = dict(month=3, day=23) # DST transition 2024
+    knowledge_func_params_dst = dict(month=3, day=23)  # DST transition 2024
     assert x_years_ago_at_date(
         event_start=datetime(2024, 3, 25, 0, tzinfo=timezone("Europe/Amsterdam")),
-        **knowledge_func_params_dst
+        **knowledge_func_params_dst,
     ) == timedelta(days=368)
 
     # Repeat test with pd.DatetimeIndex instead
-    event_start = pd.DatetimeIndex(["2024-11-19T01:00:00", "2025-11-19T02:00:00", "2022-11-19T02:00:00", "2021-11-21T03:00:00", "2021-11-21T04:00:00"], tz="utc")
+    event_start = pd.DatetimeIndex(
+        [
+            "2024-11-19T01:00:00",
+            "2025-11-19T02:00:00",
+            "2022-11-19T02:00:00",
+            "2021-11-21T03:00:00",
+            "2021-11-21T04:00:00",
+        ],
+        tz="utc",
+    )
     assert_index_equal(
-        x_years_ago_at_date(
-            event_start=event_start,
-            **knowledge_func_params
+        x_years_ago_at_date(event_start=event_start, **knowledge_func_params),
+        pd.TimedeltaIndex(
+            [
+                timedelta(days=365, hours=1),
+                timedelta(days=364, hours=2),
+                timedelta(days=364, hours=2),
+                timedelta(days=366, hours=3),
+                timedelta(days=366, hours=4),
+            ]
         ),
-        pd.TimedeltaIndex([timedelta(days=365, hours=1), timedelta(days=364, hours=2), timedelta(days=364, hours=2), timedelta(days=366, hours=3), timedelta(days=366, hours=4)]),
     )
 
-    knowledge_func_params_2_years = dict(month=11, day=20, x=2)    
+    knowledge_func_params_2_years = dict(month=11, day=20, x=2)
 
     # Check years parameter
     assert x_years_ago_at_date(
         event_start=datetime(2024, 11, 19, 1, tzinfo=utc),
-        **knowledge_func_params_2_years
-    ) == timedelta(days=2*365, hours=1) # 365 days + 366 days - 1 day
+        **knowledge_func_params_2_years,
+    ) == timedelta(
+        days=2 * 365, hours=1
+    )  # 365 days + 366 days - 1 day
 
 
 def test_dst():