Merge pull request #25 from dawidlinek/feat/calendar-source-config

feat(calendar): allow overriding lat/lon for daylight calculation

Author: dawidlinek

epftoolbox2/data/sources/calendar.py

@@ -44,6 +44,8 @@ def __init__(
         hour: Union[str, bool] = False,
         month: Union[str, bool] = False,
         daylight: bool = False,
+        lat: Optional[float] = None,
+        lon: Optional[float] = None,
         prefix: str = "",
         freq: str = "1h",
     ):
@@ -52,8 +54,8 @@ def __init__(
 
         country_info = COUNTRY_DATA[self.country]
         self.timezone = timezone or country_info["timezone"]
-        self.lat = country_info["lat"]
-        self.lon = country_info["lon"]
+        self.lat = lat if lat is not None else country_info["lat"]
+        self.lon = lon if lon is not None else country_info["lon"]
 
         self.holidays = holidays
         self.weekday = weekday

tests/data/sources/test_calendar.py

@@ -21,6 +21,21 @@ def test_init_custom_timezone(self):
         source = CalendarSource(country="PL", timezone="UTC")
         assert source.timezone == "UTC"
 
+    def test_init_default_coordinates(self):
+        source = CalendarSource(country="PL")
+        assert source.lat == COUNTRY_DATA["PL"]["lat"]
+        assert source.lon == COUNTRY_DATA["PL"]["lon"]
+
+    def test_init_override_coordinates(self):
+        source = CalendarSource(country="PL", lat=54.3520, lon=18.6466)
+        assert source.lat == 54.3520
+        assert source.lon == 18.6466
+
+    def test_init_override_lat_only(self):
+        source = CalendarSource(country="PL", lat=54.3520)
+        assert source.lat == 54.3520
+        assert source.lon == COUNTRY_DATA["PL"]["lon"]
+
     def test_init_invalid_holidays_value(self):
         with pytest.raises(ValueError, match="Invalid holidays value"):
             CalendarSource(country="PL", holidays="invalid")
@@ -168,6 +183,18 @@ def test_daylight_summer_longer(self):
         winter = source.fetch(pd.Timestamp("2024-12-21", tz="UTC"), pd.Timestamp("2024-12-22", tz="UTC"))
         assert summer["daylight_hours"].iloc[0] > winter["daylight_hours"].iloc[0]
 
+    def test_daylight_respects_coordinate_override(self):
+        # A more northern latitude has longer daylight at the summer solstice.
+        warsaw = CalendarSource(country="PL", holidays=False, weekday=False, daylight=True)
+        northern = CalendarSource(
+            country="PL", holidays=False, weekday=False, daylight=True, lat=59.3293, lon=18.0686
+        )
+        start = pd.Timestamp("2024-06-21", tz="UTC")
+        end = pd.Timestamp("2024-06-22", tz="UTC")
+        warsaw_df = warsaw.fetch(start, end)
+        northern_df = northern.fetch(start, end)
+        assert northern_df["daylight_hours"].iloc[0] > warsaw_df["daylight_hours"].iloc[0]
+
 
 class TestPrefix:
     def test_custom_prefix(self):