You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
329 lines
9.8 KiB
Python
329 lines
9.8 KiB
Python
"""
|
|
Date utility functions for quarter handling and time parameter resolution
|
|
"""
|
|
|
|
from datetime import datetime, timezone, timedelta, date
|
|
from typing import Callable, List, Tuple, Optional, Union
|
|
import re
|
|
|
|
|
|
def quarter_to_date_range(quarter: str) -> Tuple[datetime, datetime]:
|
|
"""
|
|
Convert quarter string to start and end dates
|
|
|
|
Args:
|
|
quarter: Quarter in format YYYYQN (e.g., "2020Q1")
|
|
|
|
Returns:
|
|
Tuple of (start_date, end_date) as timezone-aware datetime objects
|
|
|
|
Raises:
|
|
ValueError: If quarter format is invalid
|
|
"""
|
|
if not re.match(r'^\d{4}Q[1-4]$', quarter):
|
|
raise ValueError(f"Invalid quarter format: {quarter}. Expected format: YYYYQN (e.g., 2020Q1)")
|
|
|
|
year = int(quarter[:4])
|
|
quarter_num = int(quarter[5])
|
|
|
|
# Quarter date mappings
|
|
quarter_dates = {
|
|
1: (1, 1, 3, 31), # Q1: Jan 1 - Mar 31
|
|
2: (4, 1, 6, 30), # Q2: Apr 1 - Jun 30
|
|
3: (7, 1, 9, 30), # Q3: Jul 1 - Sep 30
|
|
4: (10, 1, 12, 31) # Q4: Oct 1 - Dec 31
|
|
}
|
|
|
|
start_month, start_day, end_month, end_day = quarter_dates[quarter_num]
|
|
|
|
start_date = datetime(year, start_month, start_day, tzinfo=timezone.utc)
|
|
end_date = datetime(year, end_month, end_day, 23, 59, 59, tzinfo=timezone.utc)
|
|
|
|
return start_date, end_date
|
|
|
|
|
|
def quarters_to_date_range(quarters: List[str]) -> Tuple[datetime, datetime]:
|
|
"""
|
|
Convert list of quarters to overall start and end dates
|
|
|
|
Args:
|
|
quarters: List of quarters in format YYYYQN
|
|
|
|
Returns:
|
|
Tuple of (earliest_start_date, latest_end_date)
|
|
|
|
Raises:
|
|
ValueError: If any quarter format is invalid
|
|
"""
|
|
if not quarters:
|
|
raise ValueError("Quarters list cannot be empty")
|
|
|
|
date_ranges = [quarter_to_date_range(q) for q in quarters]
|
|
|
|
start_dates = [dr[0] for dr in date_ranges]
|
|
end_dates = [dr[1] for dr in date_ranges]
|
|
|
|
return min(start_dates), max(end_dates)
|
|
|
|
|
|
def parse_quarter(quarter_str: str) -> Tuple[int, int]:
|
|
"""
|
|
Parse quarter string into year and quarter number
|
|
|
|
Args:
|
|
quarter_str: Quarter in format YYYYQN
|
|
|
|
Returns:
|
|
Tuple of (year, quarter_number)
|
|
|
|
Raises:
|
|
ValueError: If quarter format is invalid
|
|
"""
|
|
if not re.match(r'^\d{4}Q[1-4]$', quarter_str):
|
|
raise ValueError(f"Invalid quarter format: {quarter_str}")
|
|
|
|
year = int(quarter_str[:4])
|
|
quarter_num = int(quarter_str[5])
|
|
|
|
return year, quarter_num
|
|
|
|
|
|
def format_quarter(year: int, quarter: int) -> str:
|
|
"""
|
|
Format year and quarter number into quarter string
|
|
|
|
Args:
|
|
year: Year (e.g., 2020)
|
|
quarter: Quarter number (1-4)
|
|
|
|
Returns:
|
|
Quarter string in format YYYYQN
|
|
|
|
Raises:
|
|
ValueError: If quarter number is invalid
|
|
"""
|
|
if quarter not in [1, 2, 3, 4]:
|
|
raise ValueError(f"Quarter must be 1-4, got: {quarter}")
|
|
|
|
return f"{year}Q{quarter}"
|
|
|
|
|
|
def get_quarter_for_date(date: datetime) -> str:
|
|
"""
|
|
Get quarter string for a given date
|
|
|
|
Args:
|
|
date: Date to convert
|
|
|
|
Returns:
|
|
Quarter string in format YYYYQN
|
|
"""
|
|
year = date.year
|
|
month = date.month
|
|
|
|
if month <= 3:
|
|
quarter = 1
|
|
elif month <= 6:
|
|
quarter = 2
|
|
elif month <= 9:
|
|
quarter = 3
|
|
else:
|
|
quarter = 4
|
|
|
|
return format_quarter(year, quarter)
|
|
|
|
|
|
def expand_quarter_range(start_quarter: str, end_quarter: str) -> List[str]:
|
|
"""
|
|
Expand quarter range into list of quarters
|
|
|
|
Args:
|
|
start_quarter: Starting quarter (e.g., "2020Q1")
|
|
end_quarter: Ending quarter (e.g., "2020Q4")
|
|
|
|
Returns:
|
|
List of quarters between start and end (inclusive)
|
|
|
|
Example:
|
|
expand_quarter_range("2020Q2", "2021Q1") -> ["2020Q2", "2020Q3", "2020Q4", "2021Q1"]
|
|
"""
|
|
start_year, start_q = parse_quarter(start_quarter)
|
|
end_year, end_q = parse_quarter(end_quarter)
|
|
|
|
quarters = []
|
|
|
|
year = start_year
|
|
quarter = start_q
|
|
|
|
while year < end_year or (year == end_year and quarter <= end_q):
|
|
quarters.append(format_quarter(year, quarter))
|
|
|
|
quarter += 1
|
|
if quarter > 4:
|
|
quarter = 1
|
|
year += 1
|
|
|
|
return quarters
|
|
|
|
|
|
def parse_period(period: str) -> Tuple[datetime, datetime]:
|
|
"""
|
|
Parse period string to start and end dates
|
|
|
|
Args:
|
|
period: Period string like "1d", "2m", "3y", "5d", "1m", "2y", "max"
|
|
|
|
Returns:
|
|
Tuple of (start_date, end_date) as timezone-aware datetime objects
|
|
|
|
Raises:
|
|
ValueError: If period format is invalid
|
|
|
|
Supported formats:
|
|
- Nd: N days (e.g., "1d", "7d", "30d")
|
|
- Nm: N months (e.g., "1m", "3m", "6m")
|
|
- Ny: N years (e.g., "1y", "2y", "5y")
|
|
- max: Maximum 20 years of historical data
|
|
"""
|
|
if not period or not isinstance(period, str):
|
|
raise ValueError("Period must be a non-empty string")
|
|
|
|
# Handle "max" period - get maximum 20 years of data
|
|
if period.lower() == "max":
|
|
# For max period, limit to 20 years to prevent excessive data requests
|
|
end_date = datetime.now(timezone.utc).replace(hour=23, minute=59, second=59, microsecond=0)
|
|
start_date = end_date - timedelta(days=20 * 365.25) # 20 years
|
|
return start_date, end_date
|
|
|
|
# Match pattern: number + unit (d/m/y)
|
|
match = re.match(r'^(\d+)([dmy])$', period.lower())
|
|
if not match:
|
|
raise ValueError(f"Invalid period format: {period}. Expected format: Nd/Nm/Ny/max (e.g., 1d, 3m, 2y, max)")
|
|
|
|
amount = int(match.group(1))
|
|
unit = match.group(2)
|
|
|
|
if amount <= 0:
|
|
raise ValueError(f"Period amount must be positive, got: {amount}")
|
|
|
|
# Calculate end date (yesterday to ensure data availability)
|
|
# Since today's data might not be available, use yesterday as end date
|
|
end_date = datetime.now(timezone.utc) - timedelta(days=1)
|
|
# Set to end of day (23:59:59)
|
|
end_date = end_date.replace(hour=23, minute=59, second=59, microsecond=0)
|
|
|
|
# Calculate start date based on unit
|
|
if unit == 'd':
|
|
# Days
|
|
start_date = end_date - timedelta(days=amount)
|
|
elif unit == 'm':
|
|
# Months (approximate using 30.44 days per month)
|
|
days_in_months = amount * 30.44
|
|
start_date = end_date - timedelta(days=days_in_months)
|
|
elif unit == 'y':
|
|
# Years (approximate using 365.25 days per year)
|
|
days_in_years = amount * 365.25
|
|
start_date = end_date - timedelta(days=days_in_years)
|
|
else:
|
|
raise ValueError(f"Unsupported unit: {unit}")
|
|
|
|
return start_date, end_date
|
|
|
|
|
|
def validate_period_format(period: str) -> bool:
|
|
"""
|
|
Validate if period string has correct format
|
|
|
|
Args:
|
|
period: Period string to validate
|
|
|
|
Returns:
|
|
True if valid, False otherwise
|
|
"""
|
|
try:
|
|
parse_period(period)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def period_to_description(period: str) -> str:
|
|
"""
|
|
Convert period string to human-readable description
|
|
|
|
Args:
|
|
period: Period string like "1d", "3m", "2y", "max"
|
|
|
|
Returns:
|
|
Human-readable description
|
|
|
|
Example:
|
|
period_to_description("1d") -> "1 day"
|
|
period_to_description("3m") -> "3 months"
|
|
period_to_description("2y") -> "2 years"
|
|
period_to_description("max") -> "Maximum 20 years of data"
|
|
"""
|
|
if period.lower() == "max":
|
|
return "Maximum 20 years of data"
|
|
|
|
match = re.match(r'^(\d+)([dmy])$', period.lower())
|
|
if not match:
|
|
return period
|
|
|
|
amount = int(match.group(1))
|
|
unit = match.group(2)
|
|
|
|
unit_names = {
|
|
'd': 'day' if amount == 1 else 'days',
|
|
'm': 'month' if amount == 1 else 'months',
|
|
'y': 'year' if amount == 1 else 'years'
|
|
}
|
|
|
|
return f"{amount} {unit_names.get(unit, unit)}"
|
|
|
|
|
|
def resolve_time_parameters(
|
|
start_date: Optional[Union[date, datetime]] = None,
|
|
end_date: Optional[Union[date, datetime]] = None,
|
|
quarters: Optional[List[str]] = None,
|
|
period: Optional[str] = None,
|
|
ticker: Optional[str] = None,
|
|
ticker_max_range_fn: Optional[Callable[[str], Tuple[datetime, datetime]]] = None
|
|
) -> Tuple[datetime, datetime]:
|
|
"""
|
|
Resolve different time parameter approaches into start and end datetimes.
|
|
|
|
Args:
|
|
start_date: Start date
|
|
end_date: End date
|
|
quarters: List of quarters like ['2024Q1', '2024Q2']
|
|
period: Period string like '1d', '3m', '2y', 'max'
|
|
ticker: Ticker symbol (used for max period)
|
|
ticker_max_range_fn: Optional callback to resolve max range for a ticker
|
|
|
|
Returns:
|
|
Tuple of (start_datetime, end_datetime) with timezone info
|
|
"""
|
|
if period:
|
|
if period.lower() == "max" and ticker and ticker_max_range_fn:
|
|
return ticker_max_range_fn(ticker)
|
|
else:
|
|
start_dt, end_dt = parse_period(period)
|
|
return start_dt, end_dt
|
|
elif quarters:
|
|
start_dt, end_dt = quarters_to_date_range(quarters)
|
|
return start_dt, end_dt
|
|
elif start_date and end_date:
|
|
if isinstance(start_date, date) and not isinstance(start_date, datetime):
|
|
start_dt = datetime.combine(start_date, datetime.min.time(), timezone.utc)
|
|
else:
|
|
start_dt = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date
|
|
|
|
if isinstance(end_date, date) and not isinstance(end_date, datetime):
|
|
end_dt = datetime.combine(end_date, datetime.max.time(), timezone.utc)
|
|
else:
|
|
end_dt = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date
|
|
|
|
return start_dt, end_dt
|
|
else:
|
|
raise ValueError("One of period, quarters, or start_date/end_date must be provided") |