ADD: Python support for decimal.Decimal

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.22.0 - TBD
+
+#### Enhancements
+- Added `price_type` argument for `DBNStore.to_df` to specify if price fields should be `fixed`, `float` or `decimal.Decimal`
+
+#### Deprecations
+- Deprecated `pretty_px` argument for `DBNStore.to_df` to be removed in a future release; the default `pretty_px=True` is now equivalent to `price_type="float"` and `pretty_px=False` is now equivalent to `price_type="fixed"`
+
 ## 0.21.0 - 2023-10-11
 
 #### Enhancements

databento/common/dbnstore.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
 import abc
+import decimal
 import itertools
 import logging
+import warnings
 from collections.abc import Generator
 from collections.abc import Iterator
+from functools import partial
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, Callable, overload
+from typing import IO, TYPE_CHECKING, Any, Callable, Literal, overload
 
 import databento_dbn
 import numpy as np
@@ -44,7 +47,6 @@
 
 INT64_NULL = 9223372036854775807
 
-
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -92,37 +94,6 @@ def is_dbn(reader: IO[bytes]) -> bool:
     return reader.read(3) == b"DBN"
 
 
-def format_dataframe(
-    df: pd.DataFrame,
-    schema: Schema,
-    pretty_px: bool,
-    pretty_ts: bool,
-) -> pd.DataFrame:
-    struct = SCHEMA_STRUCT_MAP[schema]
-
-    if schema == Schema.DEFINITION:
-        for column, type_max in DEFINITION_TYPE_MAX_MAP.items():
-            if column in df.columns:
-                df[column] = df[column].where(df[column] != type_max, np.nan)
-
-    if pretty_px:
-        for px_field in struct._price_fields:
-            df[px_field] = df[px_field].replace(INT64_NULL, np.nan) / FIXED_PRICE_SCALE
-
-    if pretty_ts:
-        for ts_field in struct._timestamp_fields:
-            df[ts_field] = pd.to_datetime(df[ts_field], errors="coerce", utc=True)
-
-    for column, dtype in SCHEMA_DTYPES_MAP[schema]:
-        if dtype.startswith("S") and column not in struct._hidden_fields:
-            df[column] = df[column].str.decode("utf-8")
-
-    index_column = "ts_event" if schema.value.startswith("ohlcv") else "ts_recv"
-    df.set_index(index_column, inplace=True)
-
-    return df
-
-
 class DataSource(abc.ABC):
     """
     Abstract base class for backing DBNStore instances with data.
@@ -791,7 +762,7 @@ def to_csv(
         path: Path | str,
         pretty_px: bool = True,
         pretty_ts: bool = True,
-        map_symbols: bool | None = None,
+        map_symbols: bool = True,
         schema: Schema | str | None = None,
     ) -> None:
         """
@@ -826,8 +797,12 @@ def to_csv(
         Requires all the data to be brought up into memory to then be written.
 
         """
+        price_type: Literal["fixed", "float"] = "fixed"
+        if pretty_px is True:
+            price_type = "float"
+
         df_iter = self.to_df(
-            pretty_px=pretty_px,
+            price_type=price_type,
             pretty_ts=pretty_ts,
             map_symbols=map_symbols,
             schema=schema,
@@ -844,7 +819,8 @@ def to_csv(
     @overload
     def to_df(
         self,
-        pretty_px: bool = ...,
+        pretty_px: bool | None = ...,
+        price_type: Literal["fixed", "float", "decimal"] = ...,
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
@@ -855,7 +831,8 @@ def to_df(
     @overload
     def to_df(
         self,
-        pretty_px: bool = ...,
+        pretty_px: bool | None = ...,
+        price_type: Literal["fixed", "float", "decimal"] = ...,
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
@@ -865,7 +842,8 @@ def to_df(
 
     def to_df(
         self,
-        pretty_px: bool = True,
+        pretty_px: bool | None = None,
+        price_type: Literal["fixed", "float", "decimal"] = "float",
         pretty_ts: bool = True,
         map_symbols: bool = True,
         schema: Schema | str | None = None,
@@ -877,9 +855,15 @@ def to_df(
         Parameters
         ----------
         pretty_px : bool, default True
+            This parameter is deprecated and will be removed in a future release.
             If all price columns should be converted from `int` to `float` at
             the correct scale (using the fixed-precision scalar 1e-9). Null
             prices are replaced with NaN.
+        price_type : str, default "float"
+            The price type to use for price fields.
+            If "fixed", prices will have a type of `int` in fixed decimal format; each unit representing 1e-9 or 0.000000001.
+            If "float", prices will have a type of `float`.
+            If "decimal", prices will be instances of `decimal.Decimal`.
         pretty_ts : bool, default True
             If all timestamp columns should be converted from UNIX nanosecond
             `int` to tz-aware UTC `pd.Timestamp`.
@@ -908,6 +892,20 @@ def to_df(
             If the schema for the array cannot be determined.
 
         """
+        if pretty_px is True:
+            warnings.warn(
+                'The argument `pretty_px` is deprecated and will be removed in a future release; `price_type="float"` can be used instead.',
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        elif pretty_px is False:
+            price_type = "fixed"
+            warnings.warn(
+                'The argument `pretty_px` is deprecated and will be removed in a future release; `price_type="fixed"` can be used instead.',
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         schema = validate_maybe_enum(schema, Schema, "schema")
         if schema is None:
             if self.schema is None:
@@ -927,7 +925,7 @@ def to_df(
             schema=schema,
             count=count,
             instrument_map=self._instrument_map,
-            pretty_px=pretty_px,
+            price_type=price_type,
             pretty_ts=pretty_ts,
             map_symbols=map_symbols,
         )
@@ -966,7 +964,7 @@ def to_json(
         path: Path | str,
         pretty_px: bool = True,
         pretty_ts: bool = True,
-        map_symbols: bool | None = None,
+        map_symbols: bool = True,
         schema: Schema | str | None = None,
     ) -> None:
         """
@@ -1000,8 +998,12 @@ def to_json(
         Requires all the data to be brought up into memory to then be written.
 
         """
+        price_type: Literal["fixed", "float"] = "fixed"
+        if pretty_px is True:
+            price_type = "float"
+
         df_iter = self.to_df(
-            pretty_px=pretty_px,
+            price_type=price_type,
             pretty_ts=pretty_ts,
             map_symbols=map_symbols,
             schema=schema,
@@ -1126,37 +1128,91 @@ def __init__(
         count: int | None,
         schema: Schema,
         instrument_map: InstrumentMap,
-        pretty_px: bool = True,
+        price_type: Literal["fixed", "float", "decimal"] = "float",
         pretty_ts: bool = True,
         map_symbols: bool = True,
     ):
         self._records = records
         self._schema = schema
         self._count = count
-        self._pretty_px = pretty_px
+        self._price_type = price_type
         self._pretty_ts = pretty_ts
         self._map_symbols = map_symbols
         self._instrument_map = instrument_map
+        self._struct = SCHEMA_STRUCT_MAP[schema]
 
     def __iter__(self) -> DataFrameIterator:
         return self
 
     def __next__(self) -> pd.DataFrame:
-        df = format_dataframe(
-            pd.DataFrame(
-                next(self._records),
-                columns=SCHEMA_COLUMNS[self._schema],
-            ),
-            schema=self._schema,
-            pretty_px=self._pretty_px,
-            pretty_ts=self._pretty_ts,
+        df = pd.DataFrame(
+            next(self._records),
+            columns=SCHEMA_COLUMNS[self._schema],
         )
 
+        if self._schema == Schema.DEFINITION:
+            self._format_definition_fields(df)
+
+        self._format_hidden_fields(df)
+
+        self._format_px(df, self._price_type)
+
+        if self._pretty_ts:
+            self._format_pretty_ts(df)
+
+        self._format_set_index(df)
+
         if self._map_symbols:
-            df_index = df.index if self._pretty_ts else pd.to_datetime(df.index, utc=True)
-            dates = [ts.date() for ts in df_index]
-            df["symbol"] = [
-               self._instrument_map.resolve(inst, dates[i]) for i, inst in enumerate(df["instrument_id"])
-            ]
+            self._format_map_symbols(df)
 
         return df
+
+    def _format_definition_fields(self, df: pd.DataFrame) -> None:
+        for column, type_max in DEFINITION_TYPE_MAX_MAP.items():
+            if column in df.columns:
+                df[column] = df[column].where(df[column] != type_max, np.nan)
+
+    def _format_hidden_fields(self, df: pd.DataFrame) -> None:
+        for column, dtype in SCHEMA_DTYPES_MAP[self._schema]:
+            hidden_fields = self._struct._hidden_fields
+            if dtype.startswith("S") and column not in hidden_fields:
+                df[column] = df[column].str.decode("utf-8")
+
+    def _format_map_symbols(self, df: pd.DataFrame) -> None:
+        df_index = df.index if self._pretty_ts else pd.to_datetime(df.index, utc=True)
+        dates = [ts.date() for ts in df_index]
+        df["symbol"] = [
+            self._instrument_map.resolve(inst, dates[i])
+            for i, inst in enumerate(df["instrument_id"])
+        ]
+
+    def _format_px(
+        self,
+        df: pd.DataFrame,
+        price_type: Literal["fixed", "float", "decimal"],
+    ) -> None:
+        px_fields = self._struct._price_fields
+
+        if price_type == "decimal":
+            for field in px_fields:
+                df[field] = (
+                    df[field].replace(INT64_NULL, np.nan).apply(decimal.Decimal)
+                    / FIXED_PRICE_SCALE
+                )
+        elif price_type == "float":
+            for field in px_fields:
+                df[field] = df[field].replace(INT64_NULL, np.nan) / FIXED_PRICE_SCALE
+        else:
+            return  # do nothing
+
+    def _format_pretty_ts(self, df: pd.DataFrame) -> None:
+        for field in self._struct._timestamp_fields:
+            df[field] = df[field].apply(
+                partial(pd.to_datetime, utc=True, errors="coerce"),
+            )
+
+    def _format_set_index(self, df: pd.DataFrame) -> None:
+        index_column = (
+            "ts_event" if self._schema.value.startswith("ohlcv") else "ts_recv"
+        )
+        df.set_index(index_column, inplace=True)