Merge pull request #44 from rcholic/dtype

support for compressing dataframes
rcholic · Dec 14, 2024 · 21b15b3 · 21b15b3
2 parents 4d24fd0 + b69b983
commit 21b15b3
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 10 deletions.
diff --git a/cschwabpy/models/__init__.py b/cschwabpy/models/__init__.py
@@ -382,19 +382,24 @@ class OptionChain(JSONSerializableBaseModel):
     ]  # key: expiration:27 value:[strike: OptionContract]
 
     def to_dataframe_pairs_by_expiration(
-        self, strip_space: bool = False
+        self, strip_space: bool = False, use_compression: bool = False
     ) -> List[OptionChainDataFrames]:
         """
         List of OptionChainDataFrames by expiration.
         Each OptionChainDataFrames object contains call and put chain in dataframe format.
         @param strip_space: Whether strip spaces in option symbols.
+        @param use_compress: Whether to compress the data frame to bare minimum data types.
         """
         results: List[OptionChainDataFrames] = []
         call_map = self.break_down_option_map(
-            self.callExpDateMap, strip_space=strip_space
+            self.callExpDateMap,
+            strip_space=strip_space,
+            should_compress=use_compression,
         )
         put_map = self.break_down_option_map(
-            self.putExpDateMap, strip_space=strip_space
+            self.putExpDateMap,
+            strip_space=strip_space,
+            should_compress=use_compression,
         )
         for expiration in call_map.keys():
             call_df = call_map[expiration]
@@ -414,8 +419,13 @@ def break_down_option_map(
         self,
         optionExpMap: Mapping[str, Mapping[str, List[OptionContract]]],
         strip_space: bool = False,
+        should_compress: bool = False,
     ) -> Mapping[str, pd.DataFrame]:
-        """Whether strip spaces in option symbols."""
+        """
+        Whether strip spaces in option symbols.
+        if should_compress is True, cast the data types
+        to bare minimum to save memory usage.
+        """
         now_unix_ts = int(util.now_unix_ts())
         result: MutableMapping[str, pd.DataFrame] = {}
         for exp_date, strike_map in optionExpMap.items():
@@ -426,7 +436,6 @@ def break_down_option_map(
             all_rows = []
             strike_df = pd.DataFrame()
             for strike_str, option_contracts in strike_map.items():
-                strike = float(strike_str)
                 for option_contract in option_contracts:
                     row = option_contract.to_dataframe_row(strip_space=strip_space)
                     row.insert(0, self.underlying.mark)
@@ -436,6 +445,25 @@ def break_down_option_map(
                 strike_df.columns = OptionChain_Headers
                 # change updated_at to now_unix_ts
                 strike_df["updated_at"] = now_unix_ts
+
+            if should_compress:
+                strike_df = strike_df.astype(
+                    {
+                        "underlying_price": "float16",
+                        "strike": "float16",
+                        "lastPrice": "float16",
+                        "openInterest": "int32",
+                        "volume": "int32",
+                        "ask": "float16",
+                        "bid": "float16",
+                        "lastPrice": "float16",
+                        "gamma": "float16",
+                        "delta": "float16",
+                        "vega": "float16",
+                        "volatility": "float16",
+                    }
+                )
+
             result[expiration] = strike_df
 
         return result
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cschwabpy"
-version = "0.1.4.3"
+version = "0.1.4.4"
 description = ""
 authors = ["Tony Wang <[email protected]>"]
 readme = "README.md"

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="CSchwabPy",
-    version="0.1.4.3",
+    version="0.1.4.4",
     description="Charles Schwab Stock & Option Trade API Client for Python.",
     long_description=long_description,
     long_description_content_type="text/markdown",

diff --git a/tests/test_models.py b/tests/test_models.py
@@ -71,6 +71,7 @@ def get_mock_response(
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_market_hours(httpx_mock: HTTPXMock) -> None:
     market_hours_json = {
         "start": "2022-04-14T09:30:00-04:00",
@@ -153,7 +154,9 @@ def test_option_chain_parsing() -> None:
     assert opt_chain_result is not None
     assert opt_chain_result.status == "SUCCESS"
 
-    opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration()
+    opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration(
+        use_compression=True
+    )
     assert opt_df_pairs is not None
     for df in opt_df_pairs:
         print(df.expiration)
@@ -190,6 +193,7 @@ def test_parsing_order():
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_order(httpx_mock: HTTPXMock):
     json_mock = get_mock_response()["single_order"]
     mocked_token = mock_tokens()
@@ -242,6 +246,7 @@ async def test_get_order(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_place_order(httpx_mock: HTTPXMock):
     mocked_token = mock_tokens()
     if os.path.exists(Path(token_store.token_output_path)):
@@ -321,6 +326,7 @@ async def test_place_order(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_cancel_order(httpx_mock: HTTPXMock):
     mocked_token = mock_tokens()
     if os.path.exists(Path(token_store.token_output_path)):
@@ -359,6 +365,7 @@ async def test_cancel_order(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_order_by_id(httpx_mock: HTTPXMock):
     json_mock = get_mock_response()["filled_order"]
     mocked_token = mock_tokens()
@@ -403,6 +410,7 @@ async def test_get_order_by_id(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_single_account(httpx_mock: HTTPXMock):
     json_mock = get_mock_response()["single_account"]
     mocked_token = mock_tokens()
@@ -449,6 +457,7 @@ async def test_get_single_account(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_securities_account(httpx_mock: HTTPXMock):
     json_mock = get_mock_response()["securities_account"]  # single_account
     mocked_token = mock_tokens()
@@ -499,6 +508,7 @@ async def test_get_securities_account(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_download_option_chain(httpx_mock: HTTPXMock):
     mock_option_chain_resp = get_mock_response()
     mocked_token = mock_tokens()
@@ -527,7 +537,9 @@ async def test_download_option_chain(httpx_mock: HTTPXMock):
         assert opt_chain_result is not None
         assert opt_chain_result.status == "SUCCESS"
 
-        opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration()
+        opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration(
+            use_compression=True
+        )
         assert opt_df_pairs is not None
         for df in opt_df_pairs:
             print(df.expiration)
@@ -552,7 +564,9 @@ async def test_download_option_chain(httpx_mock: HTTPXMock):
         assert opt_chain_result2 is not None
         assert opt_chain_result2.status == "SUCCESS"
 
-        opt_df_pairs = opt_chain_result2.to_dataframe_pairs_by_expiration()
+        opt_df_pairs = opt_chain_result2.to_dataframe_pairs_by_expiration(
+            use_compression=True
+        )
         assert opt_df_pairs is not None
         for df in opt_df_pairs:
             print(df.expiration)
@@ -565,6 +579,7 @@ async def test_download_option_chain(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_option_expirations(httpx_mock: HTTPXMock):
     mock_option_chain_resp = get_mock_response()
     mocked_token = mock_tokens()
@@ -616,6 +631,7 @@ async def test_get_option_expirations(httpx_mock: HTTPXMock):
 
 
 @pytest.mark.asyncio
+@pytest.mark.httpx_mock(assert_all_requests_were_expected=False)
 async def test_get_account_numbers(httpx_mock: HTTPXMock):
     # Mock response for account numbers API
     mock_data = get_mock_response()