diff --git a/paperone/utils.py b/paperone/utils.py index 6262469..95da1f0 100644 --- a/paperone/utils.py +++ b/paperone/utils.py @@ -1,31 +1,72 @@ from datetime import datetime, timedelta from typing import List +import pandas_market_calendars as mcal + + +NYSE = mcal.get_calendar("NYSE") + def is_trading_day(date: datetime) -> bool: - return date.weekday() not in [5, 6] + valid_days = NYSE.valid_days( + start_date=date.strftime("%Y-%m-%d"), end_date=date.strftime("%Y-%m-%d") + ) + + return len(valid_days) > 0 + + +def get_trading_days(start_date: datetime, end_date: datetime) -> List[datetime]: + valid_days = NYSE.valid_days( + start_date=start_date.strftime("%Y-%m-%d"), + end_date=end_date.strftime("%Y-%m-%d"), + ) + + return [day.to_pydatetime().replace(tzinfo=None) for day in valid_days] + + +def get_last_n_trading_days(date: datetime, n: int) -> List[datetime]: + """ + Get the last N trading days before (and including) the given date. + """ + # Add buffer for weekends/holidays (n trading days ≈ n * 1.5 calendar days) + buffer_days = int(n * 2) + start_date = date - timedelta(days=buffer_days) + + # Get all trading days in the range + trading_days = get_trading_days(start_date, date) + + # Return the last n trading days + return trading_days[-n:] if len(trading_days) >= n else trading_days + + +def get_next_trading_day(date: datetime) -> datetime: + next_day = date + timedelta(days=1) + + # Search up to 10 days ahead (handles long holiday weekends) + for i in range(10): + check_date = next_day + timedelta(days=i) + if is_trading_day(check_date): + return check_date + + raise ValueError(f"No trading day found within 10 days of {date}") + + +def get_previous_trading_day(date: datetime) -> datetime: + prev_day = date - timedelta(days=1) + + # Search up to 10 days back + for i in range(10): + check_date = prev_day - timedelta(days=i) + if is_trading_day(check_date): + return check_date + + raise ValueError(f"No trading day found within 10 days before {date}") def parse_date_yyyymmdd(date_str: str) -> datetime: + """Parse date string in YYYYMMDD format to datetime.""" return datetime.strptime(date_str, "%Y%m%d") -def days_range_from(date: datetime, n: int) -> List[datetime]: - """ - Generate a list of dates going back n days from the given date (inclusive). - - Example: date=Jan 3, n=2 → [Jan 1, Jan 2, Jan 3] - - Args: - date: The target date - n: Number of days to backtrack - - Returns: - List of datetime objects from (date - n) to date (inclusive) - """ - start_date = date - timedelta(days=n) - - return [start_date + timedelta(days=i) for i in range(n + 1)] - def format_date_readable(date: datetime) -> str: + """Format datetime to readable string (e.g., 'October 15, 2025').""" return date.strftime("%B %d, %Y") -