"""Canonical DataFrame schemas for MT5 market and account datasets.""" from __future__ import annotations from enum import StrEnum from typing import TYPE_CHECKING, Final import pandas as pd from .converters import normalize_symbol, parse_timeframe from .exceptions import Mt5SchemaError if TYPE_CHECKING: from collections.abc import Iterable __all__ = [ "DEDUP_KEYS", "KNOWN_MT5_TIME_COLUMNS", "REQUIRED_COLUMNS", "TIME_COLUMNS", "DataKind", "normalize_dataframe", "normalize_time_columns", "schema_columns", "validate_schema", ] KNOWN_MT5_TIME_COLUMNS: Final[frozenset[str]] = frozenset({ "time", "time_setup", "time_setup_msc", "time_done", "time_done_msc", "time_msc", }) _TIME_COLUMN_NAMES = KNOWN_MT5_TIME_COLUMNS class DataKind(StrEnum): """Supported MT5 dataset kinds with canonical column contracts.""" rates = "rates" ticks = "ticks" orders = "orders" positions = "positions" history_orders = "history_orders" history_deals = "history_deals" REQUIRED_COLUMNS: dict[DataKind, frozenset[str]] = { DataKind.rates: frozenset({ "time", "open", "high", "low", "close", "tick_volume", "spread", "real_volume", }), DataKind.ticks: frozenset({ "time", "bid", "ask", "last", "volume", "time_msc", "flags", "volume_real", }), DataKind.orders: frozenset({ "ticket", "time_setup", "type", "state", "symbol", "volume_current", "price_open", }), DataKind.positions: frozenset({ "ticket", "time", "type", "symbol", "volume", "price_open", "price_current", "profit", }), DataKind.history_orders: frozenset({ "ticket", "time_setup", "type", "state", "symbol", "volume_initial", "price_open", }), DataKind.history_deals: frozenset({ "ticket", "order", "time", "type", "entry", "symbol", "volume", "price", "profit", }), } _OPTIONAL_TIME_COLUMNS_BY_KIND: dict[DataKind, frozenset[str]] = { DataKind.orders: frozenset({ "time_setup_msc", "time_done", "time_done_msc", }), DataKind.history_orders: frozenset({ "time_setup_msc", "time_done", "time_done_msc", }), DataKind.positions: frozenset({"time_msc"}), } TIME_COLUMNS: dict[DataKind, frozenset[str]] = { kind: (REQUIRED_COLUMNS[kind] & _TIME_COLUMN_NAMES) | _OPTIONAL_TIME_COLUMNS_BY_KIND.get(kind, frozenset()) for kind in DataKind } DEDUP_KEYS: dict[DataKind, tuple[tuple[str, ...], ...]] = { DataKind.rates: (("symbol", "timeframe", "time"), ("symbol", "time")), DataKind.ticks: (("symbol", "time_msc"), ("symbol", "time")), DataKind.history_orders: (("ticket",), ("symbol", "time", "type")), DataKind.history_deals: (("ticket",), ("symbol", "time", "type", "entry")), } def schema_columns(kind: DataKind) -> frozenset[str]: """Return required column names for a dataset kind. Args: kind: Dataset kind. Returns: Required column names for ``kind``. """ return REQUIRED_COLUMNS[kind] def validate_schema( frame: pd.DataFrame, kind: DataKind, *, extra_required: Iterable[str] | None = None, ) -> None: """Validate that a DataFrame includes required columns for a dataset kind. Args: frame: DataFrame to validate. kind: Expected dataset kind. extra_required: Additional columns that must be present (for example ``symbol`` and ``timeframe`` on stored rate history). Raises: Mt5SchemaError: If required columns are missing. """ if frame.empty and len(frame.columns) == 0: return required = set(REQUIRED_COLUMNS[kind]) if extra_required is not None: required.update(extra_required) missing = required - set(frame.columns) if missing: msg = ( f"{kind.value} schema is missing required columns: " f"{', '.join(sorted(missing))}." ) raise Mt5SchemaError(msg) def _coerce_mt5_time_column(series: pd.Series, column: str) -> pd.Series: """Coerce one MT5 time column to UTC-aware datetimes. Returns: Series with UTC-aware datetime values. """ if pd.api.types.is_datetime64_any_dtype(series): return pd.to_datetime(series, utc=True, errors="coerce") if pd.api.types.is_numeric_dtype(series): unit = "ms" if column.endswith("_msc") else "s" return pd.to_datetime(series, unit=unit, utc=True, errors="coerce") return pd.to_datetime(series, utc=True, errors="coerce") def normalize_time_columns(frame: pd.DataFrame, kind: DataKind) -> pd.DataFrame: """Coerce dataset time columns to UTC-aware datetimes when present. Any column in :data:`KNOWN_MT5_TIME_COLUMNS` that is present in ``frame`` is normalized. Numeric MT5 epoch values use seconds for ``time``, ``time_setup``, and ``time_done``, and milliseconds for ``*_msc`` columns. Args: frame: Source DataFrame from MT5 or pdmt5. kind: Dataset kind (retained for API compatibility). Returns: DataFrame copy with normalized time columns. """ del kind normalized = frame.copy() for column in normalized.columns: if column not in _TIME_COLUMN_NAMES: continue normalized[column] = _coerce_mt5_time_column(normalized[column], column) return normalized def normalize_dataframe( frame: pd.DataFrame, kind: DataKind, *, symbol: str | None = None, timeframe: int | str | None = None, sort: bool = True, ) -> pd.DataFrame: """Normalize MT5 DataFrame columns, timestamps, and storage metadata. Ensures UTC timestamps, optionally injects ``symbol`` / ``timeframe`` for storage-oriented datasets, and sorts chronologically when a ``time`` column exists. Args: frame: Source DataFrame from MT5 or pdmt5. kind: Dataset kind guiding normalization rules. symbol: Optional symbol to inject when missing. timeframe: Optional timeframe integer or name to inject for rates. sort: Whether to sort by ``time`` or ``time_msc`` when present. Returns: Normalized DataFrame copy. """ if frame.empty and len(frame.columns) == 0: return frame.copy() normalized = normalize_time_columns(frame, kind) if symbol is not None and "symbol" not in normalized.columns: normalized.insert(0, "symbol", normalize_symbol(symbol)) if timeframe is not None and kind is DataKind.rates: tf = parse_timeframe(timeframe) if "timeframe" not in normalized.columns: insert_at = 1 if "symbol" in normalized.columns else 0 normalized.insert(insert_at, "timeframe", tf) validate_schema(normalized, kind) if sort: if "time" in normalized.columns: normalized = normalized.sort_values("time", kind="stable") elif "time_msc" in normalized.columns: normalized = normalized.sort_values("time_msc", kind="stable") normalized = normalized.reset_index(drop=True) return normalized def ensure_utc_columns(frame: pd.DataFrame, columns: Iterable[str]) -> pd.DataFrame: """Return a copy with selected columns coerced to UTC datetimes. Args: frame: Source DataFrame. columns: Column names to coerce. Returns: DataFrame copy with UTC-aware datetime columns. """ normalized = frame.copy() for column in columns: if column not in normalized.columns: continue if column in _TIME_COLUMN_NAMES: normalized[column] = _coerce_mt5_time_column(normalized[column], column) else: normalized[column] = pd.to_datetime( normalized[column], utc=True, errors="coerce" ) return normalized