wehub-resource-sync
370 lines
12 KiB
Python
370 lines
12 KiB
Python
"""Shared utilities for ML4T data download scripts.
|
|
|
|
Provides:
|
|
- Standardized argument parsing with --dry-run, --data-path, --force, --verbose
|
|
- Standardized path resolution via utils.config
|
|
- YAML config helpers for dataset download scripts
|
|
- Import checking with helpful errors
|
|
- Consistent output formatting
|
|
- Atomic file writes
|
|
- Download summary reporting
|
|
- DataBento cost acknowledgment
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
_repo_root = Path(__file__).parent.parent
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Path resolution
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def resolve_data_dir(cli_arg: Path | None = None) -> Path:
|
|
"""Resolve data directory with standardized precedence.
|
|
|
|
Priority:
|
|
1. CLI argument (highest)
|
|
2. ML4T_DATA_PATH environment variable (from utils.config)
|
|
3. <repo>/data (default)
|
|
"""
|
|
if cli_arg is not None:
|
|
path = Path(cli_arg).expanduser().resolve()
|
|
print(f"Using data path (CLI): {path}")
|
|
return path
|
|
|
|
try:
|
|
from utils.config import ML4T_DATA_PATH
|
|
|
|
print(f"Using data path (ML4T_DATA_PATH): {ML4T_DATA_PATH}")
|
|
return ML4T_DATA_PATH
|
|
except (ImportError, FileNotFoundError):
|
|
pass
|
|
|
|
env_dir = os.environ.get("ML4T_DATA_PATH")
|
|
if env_dir:
|
|
path = Path(env_dir).expanduser().resolve()
|
|
print(f"Using data path (env): {path}")
|
|
return path
|
|
|
|
default_path = _repo_root / "data"
|
|
print(f"Using data path (default): {default_path}")
|
|
default_path.mkdir(parents=True, exist_ok=True)
|
|
return default_path
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Environment
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_dotenv(env_file: Path | None = None):
|
|
"""Load environment variables from .env file."""
|
|
if env_file is None:
|
|
env_file = _repo_root / ".env"
|
|
|
|
if not env_file.exists():
|
|
return
|
|
|
|
with open(env_file) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
|
|
key, value = line.split("=", 1)
|
|
key = key.strip()
|
|
value = value.strip().strip('"').strip("'")
|
|
|
|
if value and not os.getenv(key):
|
|
os.environ[key] = value
|
|
|
|
|
|
def require_env(var: str, hint: str | None = None) -> str:
|
|
"""Get required environment variable or exit with helpful message."""
|
|
value = os.getenv(var)
|
|
if not value:
|
|
print(f"ERROR: {var} not set")
|
|
if hint:
|
|
print(f" {hint}")
|
|
print(f" Add to .env file: {var}=your-value")
|
|
sys.exit(1)
|
|
assert value is not None
|
|
return value
|
|
|
|
|
|
def check_import(module: str, install_hint: str):
|
|
"""Check if module can be imported, exit with helpful message if not."""
|
|
try:
|
|
__import__(module)
|
|
except ImportError as e:
|
|
print(f"ERROR: {module} not available")
|
|
print(f" Run: {install_hint}")
|
|
print(f" ({e})")
|
|
sys.exit(1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# YAML config helpers (from book_config)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_section(config_path: str | Path, section: str) -> dict[str, Any]:
|
|
"""Load a top-level YAML section from a config file."""
|
|
import yaml
|
|
|
|
path = Path(config_path).expanduser()
|
|
with open(path) as f:
|
|
raw = yaml.safe_load(f) or {}
|
|
return raw.get(section, {})
|
|
|
|
|
|
def resolve_storage_path(data_root: Path, configured_path: str | None, fallback: str) -> Path:
|
|
"""Resolve storage path relative to the selected ML4T data root."""
|
|
raw_path = configured_path or fallback
|
|
path = Path(raw_path).expanduser()
|
|
return path if path.is_absolute() else data_root / path
|
|
|
|
|
|
def flatten_group_values(groups: dict[str, Any], values_key: str) -> list[str]:
|
|
"""Flatten grouped config values like symbols or pairs into a unique ordered list."""
|
|
values: list[str] = []
|
|
seen: set[str] = set()
|
|
|
|
for group in groups.values():
|
|
if not isinstance(group, dict):
|
|
continue
|
|
for value in group.get(values_key, []):
|
|
if value not in seen:
|
|
values.append(value)
|
|
seen.add(value)
|
|
|
|
return values
|
|
|
|
|
|
def save_dataset_profile(
|
|
df,
|
|
data_path: str | Path,
|
|
*,
|
|
source: str,
|
|
timestamp_col: str = "timestamp",
|
|
symbol_col: str | None = "symbol",
|
|
) -> Path:
|
|
"""Generate and save a dataset profile next to a data file."""
|
|
from ml4t.data.storage.data_profile import generate_profile, get_profile_path, save_profile
|
|
|
|
path = Path(data_path)
|
|
profile = generate_profile(
|
|
df, source=source, timestamp_col=timestamp_col, symbol_col=symbol_col
|
|
)
|
|
profile_path = get_profile_path(path)
|
|
save_profile(profile, profile_path)
|
|
return profile_path
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DataBento
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def patch_databento_symbology():
|
|
"""Patch databento 0.72.0 bug where insert_metadata expects 'asset' key."""
|
|
try:
|
|
import databento.common.symbology as sym
|
|
except ImportError:
|
|
return
|
|
|
|
if getattr(sym.InstrumentMap.insert_metadata, "_ml4t_patched", False):
|
|
return
|
|
|
|
_orig = sym.InstrumentMap.insert_metadata
|
|
|
|
def _patched(self, metadata):
|
|
mappings = metadata.mappings
|
|
if mappings:
|
|
for _symbol_in, entries in mappings.items():
|
|
for entry in entries:
|
|
if "asset" not in entry and "symbol" in entry:
|
|
entry["asset"] = entry["symbol"]
|
|
|
|
class _PatchedMeta:
|
|
"""Thin wrapper that returns fixed mappings."""
|
|
|
|
def __init__(self, orig, fixed_mappings):
|
|
self._orig = orig
|
|
self._fixed = fixed_mappings
|
|
|
|
@property
|
|
def mappings(self):
|
|
return self._fixed
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._orig, name)
|
|
|
|
return _orig(self, _PatchedMeta(metadata, mappings))
|
|
|
|
_patched._ml4t_patched = True
|
|
sym.InstrumentMap.insert_metadata = _patched
|
|
|
|
|
|
DATABENTO_WARNING = """
|
|
================================================================================
|
|
DATABENTO API - PAID SERVICE WARNING
|
|
================================================================================
|
|
|
|
This download uses the DataBento API which is a PAID service.
|
|
|
|
IMPORTANT INFORMATION:
|
|
- DataBento requires registration with a credit card on file
|
|
- As of February 2026, DataBento offers a $125 sign-up credit for new accounts
|
|
- This credit is sufficient to cover the ML4T book datasets:
|
|
* CME Futures (continuous contracts): ~$75
|
|
* MBO Tick Data (3 symbols, 10 days): ~$10-15
|
|
|
|
WARNINGS:
|
|
- If you have already used your sign-up credit on other downloads,
|
|
this download WILL BE CHARGED to your credit card
|
|
- DataBento may change or remove the sign-up credit at any time
|
|
- Cost estimates are approximate; actual costs may vary
|
|
- YOU are responsible for managing your DataBento account and costs
|
|
|
|
ESTIMATED COST FOR THIS DOWNLOAD: ${estimated_cost:.2f}
|
|
|
|
================================================================================
|
|
"""
|
|
|
|
|
|
def databento_acknowledge(estimated_cost: float, force: bool = False) -> bool:
|
|
"""Display DataBento cost warning and require explicit acknowledgment."""
|
|
print(DATABENTO_WARNING.format(estimated_cost=estimated_cost))
|
|
|
|
if force:
|
|
print("--force flag set: Proceeding without interactive confirmation.")
|
|
print("By using --force, you acknowledge the above warnings.")
|
|
return True
|
|
|
|
print("To proceed with this download, type exactly: I UNDERSTAND")
|
|
print("To cancel, press Ctrl+C or type anything else.")
|
|
print()
|
|
|
|
try:
|
|
response = input("Your response: ").strip()
|
|
except (KeyboardInterrupt, EOFError):
|
|
print("\nDownload cancelled.")
|
|
return False
|
|
|
|
if response == "I UNDERSTAND":
|
|
print("\nAcknowledgment received. Proceeding with download...")
|
|
return True
|
|
else:
|
|
print(f"\nResponse '{response}' does not match 'I UNDERSTAND'.")
|
|
print("Download cancelled for your protection.")
|
|
return False
|
|
|
|
|
|
def databento_estimate_only_notice(estimated_cost: float) -> None:
|
|
"""Print cost estimate without prompting for download."""
|
|
print("\n" + "=" * 70)
|
|
print("DATABENTO COST ESTIMATE (No download - estimate only)")
|
|
print("=" * 70)
|
|
print(f"\n Estimated cost: ${estimated_cost:.2f}")
|
|
print()
|
|
print(" Note: As of Feb 2026, new DataBento accounts receive $125 credit.")
|
|
print(" If your credit is exhausted, this amount will be charged.")
|
|
print()
|
|
print(" To proceed with download, run without --estimate-only flag.")
|
|
print("=" * 70 + "\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Output formatting
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def print_section(title: str, char: str = "=", width: int = 60):
|
|
"""Print a formatted section header."""
|
|
print("\n" + char * width)
|
|
print(title)
|
|
print(char * width)
|
|
|
|
|
|
def atomic_write_parquet(df, path: Path):
|
|
"""Write Polars DataFrame to parquet with atomic rename."""
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
tmp_file = path.parent / f".{path.name}.tmp"
|
|
|
|
try:
|
|
df.write_parquet(tmp_file)
|
|
tmp_file.replace(path)
|
|
except Exception as e:
|
|
if tmp_file.exists():
|
|
tmp_file.unlink()
|
|
raise e
|
|
|
|
|
|
def get_repo_root() -> Path:
|
|
"""Get repository root directory."""
|
|
return _repo_root
|
|
|
|
|
|
def create_base_parser(description: str) -> argparse.ArgumentParser:
|
|
"""Create argument parser with standard ML4T download flags."""
|
|
parser = argparse.ArgumentParser(
|
|
description=description,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
parser.add_argument(
|
|
"--data-path",
|
|
type=Path,
|
|
default=None,
|
|
help="Data storage location (default: $ML4T_DATA_PATH or repo/data)",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Show what would be done without downloading",
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Force re-download even if data exists",
|
|
)
|
|
parser.add_argument(
|
|
"--verbose",
|
|
"-v",
|
|
action="store_true",
|
|
help="Verbose output",
|
|
)
|
|
return parser
|
|
|
|
|
|
def print_download_summary(stats: dict, dry_run: bool = False) -> None:
|
|
"""Print standardized download summary."""
|
|
prefix = "[DRY RUN] " if dry_run else ""
|
|
print_section(f"{prefix}SUMMARY")
|
|
for key, value in stats.items():
|
|
display_key = key.replace("_", " ").title()
|
|
if isinstance(value, int) and value > 1000:
|
|
print(f" {display_key}: {value:,}")
|
|
elif isinstance(value, float):
|
|
print(f" {display_key}: {value:.2f}")
|
|
else:
|
|
print(f" {display_key}: {value}")
|
|
|
|
|
|
def print_dry_run_notice() -> None:
|
|
"""Print notice that this is a dry run."""
|
|
print("\n" + "=" * 60)
|
|
print("DRY RUN - No data will be downloaded")
|
|
print("Remove --dry-run to actually download")
|
|
print("=" * 60 + "\n")
|