brighter-trading/src/utils.py

87 lines
2.7 KiB
Python

"""
Utility functions for BrighterTrading.
"""
import datetime
import math
import numpy as np
import pandas as pd
def sanitize_for_json(value):
"""
Convert a value to a JSON-serializable type.
Handles numpy types, pandas types, datetime, inf, nan, and nested structures.
:param value: Any value that needs to be JSON serialized
:return: JSON-serializable value
"""
# Handle None
if value is None:
return None
# Handle pandas NA/NaT (safely check for scalar values only)
try:
if pd.isna(value) and not isinstance(value, (dict, list, tuple, set)):
return None
except (TypeError, ValueError):
pass # pd.isna fails on some types, that's fine
# Handle numpy integer types
if isinstance(value, np.integer):
return int(value)
# Handle numpy floating types
elif isinstance(value, np.floating):
float_val = float(value)
if math.isinf(float_val) or math.isnan(float_val):
return None
return float_val
# Handle numpy boolean
elif isinstance(value, np.bool_):
return bool(value)
# Handle numpy arrays
elif isinstance(value, np.ndarray):
return [sanitize_for_json(v) for v in value.tolist()]
# Handle pandas Timestamp
elif isinstance(value, pd.Timestamp):
return value.isoformat()
# Handle datetime objects
elif isinstance(value, (datetime.datetime, datetime.date)):
return value.isoformat()
# Handle pandas Timedelta
elif isinstance(value, pd.Timedelta):
return str(value)
# Handle regular floats with inf/nan
elif isinstance(value, float):
if math.isinf(value) or math.isnan(value):
return None
return value
# Handle bytes
elif isinstance(value, bytes):
try:
return value.decode('utf-8')
except UnicodeDecodeError:
return value.hex()
# Handle nested dicts
elif isinstance(value, dict):
return {str(k): sanitize_for_json(v) for k, v in value.items()}
# Handle lists
elif isinstance(value, list):
return [sanitize_for_json(v) for v in value]
# Handle tuples
elif isinstance(value, tuple):
return tuple(sanitize_for_json(v) for v in value)
# Handle sets
elif isinstance(value, set):
return [sanitize_for_json(v) for v in value]
# Return other types as-is (str, int, bool, etc.)
return value
def sanitize_dict_for_json(data: dict) -> dict:
"""
Sanitize a dictionary for JSON serialization.
Convenience wrapper around sanitize_for_json for dictionaries.
:param data: Dictionary to sanitize
:return: JSON-serializable dictionary
"""
return sanitize_for_json(data)