MOD: Add PriceType enum for validating price_type

databento · Oct 8, 2024 · bd0db0f · bd0db0f
1 parent 916e060
commit bd0db0f
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
 This release drops support for Python 3.8 which has reached end-of-life.
 
 #### Enhancements
+- Added `PriceType` enum for validation of `price_type` parameter in `DBNStore.to_df`
 - Upgraded `databento-dbn` to 0.22.1
 
 #### Bug fixes

diff --git a/databento/common/dbnstore.py b/databento/common/dbnstore.py
@@ -47,6 +47,7 @@
 from databento.common.constants import DEFINITION_TYPE_MAX_MAP
 from databento.common.constants import SCHEMA_STRUCT_MAP
 from databento.common.constants import SCHEMA_STRUCT_MAP_V1
+from databento.common.enums import PriceType
 from databento.common.error import BentoError
 from databento.common.error import BentoWarning
 from databento.common.symbology import InstrumentMap
@@ -848,7 +849,7 @@ def to_csv(
     @overload
     def to_df(
         self,
-        price_type: Literal["fixed", "float", "decimal"] = ...,
+        price_type: PriceType | str = ...,
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
@@ -859,7 +860,7 @@ def to_df(
     @overload
     def to_df(
         self,
-        price_type: Literal["fixed", "float", "decimal"] = ...,
+        price_type: PriceType | str = ...,
         pretty_ts: bool = ...,
         map_symbols: bool = ...,
         schema: Schema | str | None = ...,
@@ -869,7 +870,7 @@ def to_df(
 
     def to_df(
         self,
-        price_type: Literal["fixed", "float", "decimal"] = "float",
+        price_type: PriceType | str = PriceType.FLOAT,
         pretty_ts: bool = True,
         map_symbols: bool = True,
         schema: Schema | str | None = None,
@@ -883,7 +884,7 @@ def to_df(
 
         Parameters
         ----------
-        price_type : str, default "float"
+        price_type : PriceType or 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`.
@@ -918,6 +919,7 @@ def to_df(
             If the DBN schema is unspecified and cannot be determined.
 
         """
+        price_type = validate_enum(price_type, PriceType, "price_type")
         schema = validate_maybe_enum(schema, Schema, "schema")
 
         if isinstance(tz, Default):
@@ -1422,7 +1424,7 @@ def __init__(
         struct_type: type[DBNRecord],
         instrument_map: InstrumentMap,
         tz: pytz.BaseTzInfo,
-        price_type: Literal["fixed", "float", "decimal"] = "float",
+        price_type: PriceType = PriceType.FLOAT,
         pretty_ts: bool = True,
         map_symbols: bool = True,
     ):
@@ -1499,16 +1501,16 @@ def _format_timezone(self, df: pd.DataFrame) -> None:
     def _format_px(
         self,
         df: pd.DataFrame,
-        price_type: Literal["fixed", "float", "decimal"],
+        price_type: PriceType,
     ) -> None:
         px_fields = self._struct_type._price_fields
 
-        if price_type == "decimal":
+        if price_type == PriceType.DECIMAL:
             df[px_fields] = (
                 df[px_fields].replace(UNDEF_PRICE, np.nan).applymap(decimal.Decimal)
                 / FIXED_PRICE_SCALE
             )
-        elif price_type == "float":
+        elif price_type == PriceType.FLOAT:
             df[px_fields] = df[px_fields].replace(UNDEF_PRICE, np.nan) / FIXED_PRICE_SCALE
         else:
             return  # do nothing

diff --git a/databento/common/enums.py b/databento/common/enums.py
@@ -225,5 +225,21 @@ class RecordFlags(StringyMixin, IntFlag):  # type: ignore
 @unique
 @coercible
 class ReconnectPolicy(StringyMixin, str, Enum):
+    """
+    Live session reconnection policy.
+    """
+
     NONE = "none"
     RECONNECT = "reconnect"
+
+
+@unique
+@coercible
+class PriceType(StringyMixin, str, Enum):
+    """
+    Price type for DataFrame price fields.
+    """
+
+    FIXED = "fixed"
+    FLOAT = "float"
+    DECIMAL = "decimal"
diff --git a/tests/test_historical_bento.py b/tests/test_historical_bento.py
@@ -646,6 +646,18 @@ def test_to_df_with_price_type_handles_null(
     assert all(np.isnan(df_pretty["strike_price"]))
 
 
+def test_to_df_with_price_type_invalid(
+    test_data: Callable[[Dataset, Schema], bytes],
+) -> None:
+    # Arrange
+    stub_data = test_data(Dataset.GLBX_MDP3, Schema.DEFINITION)
+    data = DBNStore.from_bytes(data=stub_data)
+
+    # Act, Assert
+    with pytest.raises(ValueError):
+        data.to_df(price_type="US/Eastern")
+
+
 @pytest.mark.parametrize(
     "dataset",
     [