diff --git a/pyfacts/pyfacts.py b/pyfacts/pyfacts.py index 8949925..fdeff56 100644 --- a/pyfacts/pyfacts.py +++ b/pyfacts/pyfacts.py @@ -76,32 +76,28 @@ def create_date_series( 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}") - 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 = [] - - for i in range(0, int(datediff)): - diff = {frequency.freq_type: frequency.value * i} + counter = 0 + while counter < 100000: + diff = {frequency.freq_type: frequency.value * counter} date = start_date + relativedelta(**diff) if eomonth: - replacement = {"month": date.month + 1} if date.month < 12 else {"year": date.year + 1, "month": 1} - date = date.replace(day=1).replace(**replacement) - relativedelta(days=1) + date += relativedelta(months=1, day=1, days=-1) - if date <= end_date: - if frequency.days > 1 or not skip_weekends: - dates.append(date) - elif date.weekday() < 5: - dates.append(date) + if date > end_date: + if not ensure_coverage: + break + elif dates[-1] >= end_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")