Add async_get_empty_df_with_datetime_index

pypmanager/helpers/chart.py

@@ -17,6 +17,7 @@
 )
 from pypmanager.ingest.transaction.transaction_registry import TransactionRegistry
 from pypmanager.settings import Settings
+from pypmanager.utils.dt import async_get_empty_df_with_datetime_index
 
 
 @strawberry.type
@@ -70,15 +71,7 @@ async def async_get_market_data_and_transaction(
     # Create a date range between start_date and end_date, excluding weekends
     # Set start date to 1980-01-01 to ensure all dates are included
     # Convert the date range to a DataFrame and set index
-    df_date_range = pd.DataFrame(
-        pd.date_range(
-            start=pd.Timestamp("1980-01-01"),
-            end=pd.Timestamp("2030-12-31"),
-            freq="B",
-            tz=Settings.system_time_zone,
-        ),
-        columns=["date"],
-    ).set_index("date")
+    df_date_range = await async_get_empty_df_with_datetime_index()
 
     # Merge the date range DataFrame with df_transactions on the date index
     df_transaction_with_date = df_date_range.join(

pypmanager/utils/dt.py

@@ -1,8 +1,12 @@
 """Date utilities."""
 
-from datetime import datetime, timedelta
+from __future__ import annotations
+
+from datetime import date, datetime, timedelta
 from enum import IntEnum
 
+import pandas as pd
+
 from pypmanager.settings import Settings
 
 QUARTER_ENDS = [(3, 31), (6, 30), (9, 30), (12, 31)]
@@ -65,3 +69,29 @@ async def async_get_last_n_quarters(no_quarters: int) -> list[datetime]:
         quarter_ends.append(last_quarter)
 
     return quarter_ends[::-1]
+
+
+async def async_get_empty_df_with_datetime_index(
+    start_date: date | None = None,
+    end_date: date | None = None,
+) -> pd.DataFrame:
+    """
+    Return an empty DataFrame with a DatetimeIndex.
+
+    Only weekdays are included in the date range.
+    """
+    if start_date is None:
+        start_date = date(1980, 1, 1)
+
+    if end_date is None:
+        end_date = date(2030, 12, 31)
+
+    return pd.DataFrame(
+        pd.date_range(
+            start=pd.Timestamp(start_date),
+            end=pd.Timestamp(end_date),
+            freq="B",
+            tz=Settings.system_time_zone,
+        ),
+        columns=["date"],
+    ).set_index("date")

tests/utils/test_dt.py

@@ -6,7 +6,11 @@
 import pytest
 
 from pypmanager.settings import Settings
-from pypmanager.utils.dt import async_get_last_n_quarters, get_previous_quarter
+from pypmanager.utils.dt import (
+    async_get_empty_df_with_datetime_index,
+    async_get_last_n_quarters,
+    get_previous_quarter,
+)
 
 
 @pytest.mark.asyncio
@@ -46,3 +50,12 @@ async def test_async_get_last_n_quarters(
 def test_get_previous_quarter(report_date: datetime, expected_date: datetime) -> None:
     """Test function get_previous_quarter."""
     assert get_previous_quarter(report_date=report_date) == expected_date
+
+
+@pytest.mark.asyncio
+async def test_async_get_empty_df_with_datetime_index() -> None:
+    """Test function async_get_empty_df_with_datetime_index."""
+    result = await async_get_empty_df_with_datetime_index()
+    assert len(result) == 13306
+    assert result.index[0] == datetime(1980, 1, 1, tzinfo=Settings.system_time_zone)
+    assert result.index[-1] == datetime(2030, 12, 31, tzinfo=Settings.system_time_zone)