Compare commits

...

3 Commits

2 changed files with 21 additions and 25 deletions

View File

@ -76,32 +76,28 @@ def create_date_series(
if eomonth and frequency.days < AllFrequencies.M.days: if eomonth and frequency.days < AllFrequencies.M.days:
raise ValueError(f"eomonth cannot be set to True if frequency is higher than {AllFrequencies.M.name}") raise ValueError(f"eomonth cannot be set to True if frequency is higher than {AllFrequencies.M.name}")
if ensure_coverage:
if frequency.days == 1 and skip_weekends and end_date.weekday() > 4:
extend_by_days = 7 - end_date.weekday()
end_date += relativedelta(days=extend_by_days)
else:
end_date += relativedelta(days=frequency.days)
# TODO: Add code to ensure coverage for other frequencies as well
datediff = (end_date - start_date).days / frequency.days + 1
dates = [] dates = []
counter = 0
for i in range(0, int(datediff)): while counter < 100000:
diff = {frequency.freq_type: frequency.value * i} diff = {frequency.freq_type: frequency.value * counter}
date = start_date + relativedelta(**diff) date = start_date + relativedelta(**diff)
if eomonth: if eomonth:
replacement = {"month": date.month + 1} if date.month < 12 else {"year": date.year + 1, "month": 1} date += relativedelta(months=1, day=1, days=-1)
date = date.replace(day=1).replace(**replacement) - relativedelta(days=1)
if date <= end_date: if date > end_date:
if frequency.days > 1 or not skip_weekends: if not ensure_coverage:
dates.append(date) break
elif date.weekday() < 5: elif dates[-1] >= end_date:
dates.append(date) break
counter += 1
if frequency.days > 1 or not skip_weekends:
dates.append(date)
elif date.weekday() < 5:
dates.append(date)
else:
raise ValueError("Cannot generate a series containing more than 100000 dates")
return Series(dates, dtype="date") return Series(dates, dtype="date")

View File

@ -30,7 +30,7 @@ class TestDateSeries:
def test_monthly(self): def test_monthly(self):
start_date = datetime.datetime(2020, 1, 1) start_date = datetime.datetime(2020, 1, 1)
end_date = datetime.datetime(2020, 12, 31) end_date = datetime.datetime(2020, 12, 31)
d = create_date_series(start_date, end_date, frequency="M") d = create_date_series(start_date, end_date, frequency="M", ensure_coverage=False)
assert len(d) == 12 assert len(d) == 12
d = create_date_series(start_date, end_date, frequency="M", eomonth=True) d = create_date_series(start_date, end_date, frequency="M", eomonth=True)
@ -327,7 +327,7 @@ class TestExpand:
ts_data = create_test_data(AllFrequencies.M, num=6) ts_data = create_test_data(AllFrequencies.M, num=6)
ts = TimeSeries(ts_data, "M") ts = TimeSeries(ts_data, "M")
expanded_ts = ts.expand("W", "ffill") expanded_ts = ts.expand("W", "ffill")
assert len(expanded_ts) == 22 assert len(expanded_ts) == 23
assert expanded_ts.frequency.name == "weekly" assert expanded_ts.frequency.name == "weekly"
assert expanded_ts.iloc[0][1] == expanded_ts.iloc[1][1] assert expanded_ts.iloc[0][1] == expanded_ts.iloc[1][1]
@ -346,7 +346,7 @@ class TestShrink:
ts = TimeSeries(ts_data, "D") ts = TimeSeries(ts_data, "D")
shrunk_ts_w = ts.shrink("W", "ffill") shrunk_ts_w = ts.shrink("W", "ffill")
shrunk_ts_m = ts.shrink("M", "ffill") shrunk_ts_m = ts.shrink("M", "ffill")
assert len(shrunk_ts_w) == 143 assert len(shrunk_ts_w) == 144
assert len(shrunk_ts_m) == 34 assert len(shrunk_ts_m) == 34
def test_weekly_to_smaller(self, create_test_data): def test_weekly_to_smaller(self, create_test_data):
@ -376,7 +376,7 @@ class TestTransform:
ts = TimeSeries(ts_data, "D") ts = TimeSeries(ts_data, "D")
tst = ts.transform("W", "mean") tst = ts.transform("W", "mean")
assert isinstance(tst, TimeSeries) assert isinstance(tst, TimeSeries)
assert len(tst) == 156 assert len(tst) == 157
assert "2017-01-30" in tst assert "2017-01-30" in tst
assert tst.iloc[4] == (datetime.datetime(2017, 1, 30), 1021.19) assert tst.iloc[4] == (datetime.datetime(2017, 1, 30), 1021.19)