271 lines
6.4 KiB
Python
271 lines
6.4 KiB
Python
"""
|
|
ML-Specific DateTime Utilities
|
|
|
|
DateTime utilities for machine learning operations, specifically for:
|
|
- Prophet forecasting model (requires timezone-naive datetimes)
|
|
- Pandas DataFrame datetime operations
|
|
- Time series data processing
|
|
"""
|
|
|
|
from datetime import datetime, timezone
|
|
from typing import Union
|
|
import pandas as pd
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def ensure_timezone_aware(dt: datetime, default_tz=timezone.utc) -> datetime:
|
|
"""
|
|
Ensure a datetime is timezone-aware.
|
|
|
|
Args:
|
|
dt: Datetime to check
|
|
default_tz: Timezone to apply if datetime is naive (default: UTC)
|
|
|
|
Returns:
|
|
Timezone-aware datetime
|
|
"""
|
|
if dt is None:
|
|
return None
|
|
|
|
if dt.tzinfo is None:
|
|
return dt.replace(tzinfo=default_tz)
|
|
return dt
|
|
|
|
|
|
def ensure_timezone_naive(dt: datetime) -> datetime:
|
|
"""
|
|
Remove timezone information from a datetime.
|
|
|
|
Args:
|
|
dt: Datetime to process
|
|
|
|
Returns:
|
|
Timezone-naive datetime
|
|
"""
|
|
if dt is None:
|
|
return None
|
|
|
|
if dt.tzinfo is not None:
|
|
return dt.replace(tzinfo=None)
|
|
return dt
|
|
|
|
|
|
def normalize_datetime_to_utc(dt: Union[datetime, pd.Timestamp]) -> datetime:
|
|
"""
|
|
Normalize any datetime to UTC timezone-aware datetime.
|
|
|
|
Args:
|
|
dt: Datetime or pandas Timestamp to normalize
|
|
|
|
Returns:
|
|
UTC timezone-aware datetime
|
|
"""
|
|
if dt is None:
|
|
return None
|
|
|
|
if isinstance(dt, pd.Timestamp):
|
|
dt = dt.to_pydatetime()
|
|
|
|
if dt.tzinfo is None:
|
|
return dt.replace(tzinfo=timezone.utc)
|
|
|
|
return dt.astimezone(timezone.utc)
|
|
|
|
|
|
def normalize_dataframe_datetime_column(
|
|
df: pd.DataFrame,
|
|
column: str,
|
|
target_format: str = 'naive'
|
|
) -> pd.DataFrame:
|
|
"""
|
|
Normalize a datetime column in a dataframe to consistent format.
|
|
|
|
Args:
|
|
df: DataFrame to process
|
|
column: Name of datetime column
|
|
target_format: 'naive' or 'aware' (UTC)
|
|
|
|
Returns:
|
|
DataFrame with normalized datetime column
|
|
"""
|
|
if column not in df.columns:
|
|
logger.warning(f"Column {column} not found in dataframe")
|
|
return df
|
|
|
|
df[column] = pd.to_datetime(df[column])
|
|
|
|
if target_format == 'naive':
|
|
if df[column].dt.tz is not None:
|
|
df[column] = df[column].dt.tz_localize(None)
|
|
elif target_format == 'aware':
|
|
if df[column].dt.tz is None:
|
|
df[column] = df[column].dt.tz_localize(timezone.utc)
|
|
else:
|
|
df[column] = df[column].dt.tz_convert(timezone.utc)
|
|
else:
|
|
raise ValueError(f"Invalid target_format: {target_format}. Must be 'naive' or 'aware'")
|
|
|
|
return df
|
|
|
|
|
|
def prepare_prophet_datetime(df: pd.DataFrame, datetime_col: str = 'ds') -> pd.DataFrame:
|
|
"""
|
|
Prepare datetime column for Prophet (requires timezone-naive datetimes).
|
|
|
|
Args:
|
|
df: DataFrame with datetime column
|
|
datetime_col: Name of datetime column (default: 'ds')
|
|
|
|
Returns:
|
|
DataFrame with Prophet-compatible datetime column
|
|
"""
|
|
df = df.copy()
|
|
df = normalize_dataframe_datetime_column(df, datetime_col, target_format='naive')
|
|
return df
|
|
|
|
|
|
def safe_datetime_comparison(dt1: datetime, dt2: datetime) -> int:
|
|
"""
|
|
Safely compare two datetimes, handling timezone mismatches.
|
|
|
|
Args:
|
|
dt1: First datetime
|
|
dt2: Second datetime
|
|
|
|
Returns:
|
|
-1 if dt1 < dt2, 0 if equal, 1 if dt1 > dt2
|
|
"""
|
|
dt1_utc = normalize_datetime_to_utc(dt1)
|
|
dt2_utc = normalize_datetime_to_utc(dt2)
|
|
|
|
if dt1_utc < dt2_utc:
|
|
return -1
|
|
elif dt1_utc > dt2_utc:
|
|
return 1
|
|
else:
|
|
return 0
|
|
|
|
|
|
def get_current_utc() -> datetime:
|
|
"""
|
|
Get current datetime in UTC with timezone awareness.
|
|
|
|
Returns:
|
|
Current UTC datetime
|
|
"""
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def convert_timestamp_to_datetime(timestamp: Union[int, float, str]) -> datetime:
|
|
"""
|
|
Convert various timestamp formats to datetime.
|
|
|
|
Args:
|
|
timestamp: Unix timestamp (seconds or milliseconds) or ISO string
|
|
|
|
Returns:
|
|
UTC timezone-aware datetime
|
|
"""
|
|
if isinstance(timestamp, str):
|
|
dt = pd.to_datetime(timestamp)
|
|
return normalize_datetime_to_utc(dt)
|
|
|
|
if timestamp > 1e10:
|
|
timestamp = timestamp / 1000
|
|
|
|
dt = datetime.fromtimestamp(timestamp, tz=timezone.utc)
|
|
return dt
|
|
|
|
|
|
def align_dataframe_dates(
|
|
dfs: list[pd.DataFrame],
|
|
date_column: str = 'ds',
|
|
method: str = 'inner'
|
|
) -> list[pd.DataFrame]:
|
|
"""
|
|
Align multiple dataframes to have the same date range.
|
|
|
|
Args:
|
|
dfs: List of DataFrames to align
|
|
date_column: Name of the date column
|
|
method: 'inner' (intersection) or 'outer' (union)
|
|
|
|
Returns:
|
|
List of aligned DataFrames
|
|
"""
|
|
if not dfs:
|
|
return []
|
|
|
|
if len(dfs) == 1:
|
|
return dfs
|
|
|
|
all_dates = None
|
|
|
|
for df in dfs:
|
|
if date_column not in df.columns:
|
|
continue
|
|
|
|
dates = set(pd.to_datetime(df[date_column]).dt.date)
|
|
|
|
if all_dates is None:
|
|
all_dates = dates
|
|
else:
|
|
if method == 'inner':
|
|
all_dates = all_dates.intersection(dates)
|
|
elif method == 'outer':
|
|
all_dates = all_dates.union(dates)
|
|
|
|
aligned_dfs = []
|
|
for df in dfs:
|
|
if date_column not in df.columns:
|
|
aligned_dfs.append(df)
|
|
continue
|
|
|
|
df = df.copy()
|
|
df[date_column] = pd.to_datetime(df[date_column])
|
|
df['_date_only'] = df[date_column].dt.date
|
|
df = df[df['_date_only'].isin(all_dates)]
|
|
df = df.drop('_date_only', axis=1)
|
|
aligned_dfs.append(df)
|
|
|
|
return aligned_dfs
|
|
|
|
|
|
def fill_missing_dates(
|
|
df: pd.DataFrame,
|
|
date_column: str = 'ds',
|
|
freq: str = 'D',
|
|
fill_value: float = 0.0
|
|
) -> pd.DataFrame:
|
|
"""
|
|
Fill missing dates in a DataFrame with a specified frequency.
|
|
|
|
Args:
|
|
df: DataFrame with date column
|
|
date_column: Name of the date column
|
|
freq: Pandas frequency string ('D' for daily, 'H' for hourly, etc.)
|
|
fill_value: Value to fill for missing dates
|
|
|
|
Returns:
|
|
DataFrame with filled dates
|
|
"""
|
|
df = df.copy()
|
|
df[date_column] = pd.to_datetime(df[date_column])
|
|
|
|
df = df.set_index(date_column)
|
|
|
|
full_range = pd.date_range(
|
|
start=df.index.min(),
|
|
end=df.index.max(),
|
|
freq=freq
|
|
)
|
|
|
|
df = df.reindex(full_range, fill_value=fill_value)
|
|
|
|
df = df.reset_index()
|
|
df = df.rename(columns={'index': date_column})
|
|
|
|
return df
|