Merge pull request #4 from rcholic/modeling4

udated test and option chain download

README.md

@@ -47,6 +47,16 @@ schwab_client.get_tokens_manually()
 from_date = 2024-07-01
 to_date = 2024-07-01
 ticker = '$SPX'
-asyncio.run(schwab_client.download_option_chain(ticker, from_date, to_date))
+asyncio.run(opt_chain_result = schwab_client.download_option_chain(ticker, from_date, to_date))
+
+# get call-put dataframe pairs by expiration
+opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration()
+
+for df in opt_df_pairs:
+    print(df.expiration)
+    print(f"call dataframe size: {df.call_df.shape}. expiration: {df.expiration}")
+    print(f"put dataframe size: {df.put_df.shape}. expiration: {df.expiration}")
+    print(df.call_df.head(5))
+    print(df.put_df.head(5))
 
 ```

cschwabpy/SchwabAsyncClient.py

@@ -1,5 +1,5 @@
 from cschwabpy.models.token import Tokens, ITokenStore, LocalTokenStore
-from cschwabpy.models import OptionChainQueryFilter, OptionContractType
+from cschwabpy.models import OptionChainQueryFilter, OptionContractType, OptionChain
 from typing import Optional, List, Mapping
 from cschwabpy.costants import (
     SCHWAB_API_BASE_URL,
@@ -100,7 +100,9 @@ async def download_option_chain(
         from_date: str,
         to_date: str,
         contract_type: str = "ALL",
-    ) -> None:
+    ) -> OptionChain:
+        await self._ensure_valid_access_token()
+
         query_filter = OptionChainQueryFilter(
             symbol=underlying_symbol,
             contractType=OptionContractType(contract_type),
@@ -110,14 +112,14 @@ async def download_option_chain(
         target_url = (
             f"{SCHWAB_MARKET_DATA_API_BASE_URL}/chains?{query_filter.to_query_params()}"
         )
-        print("target_url: ", target_url)
-        print("auth header: ", self.__auth_header())
+
         client = httpx.AsyncClient() if self.__client is None else self.__client
         try:
             response = await client.get(
                 url=target_url, params={}, headers=self.__auth_header()
             )
             json_res = response.json()
+            return OptionChain(**json_res)
         finally:
             if not self.__keep_client_alive:
                 await client.aclose()

cschwabpy/models/__init__.py

@@ -1,9 +1,30 @@
 """models folder."""
+from datetime import datetime
+from dataclasses import dataclass
 from pydantic import BaseModel, ConfigDict, Field
 from typing import MutableMapping, Mapping, MutableSet, Any, List, Optional
 from enum import Enum
+import cschwabpy.util as util
 import pandas as pd
 
+OptionChain_Headers = [
+    "underlying_price",
+    "strike",
+    "symbol",
+    "last_price",
+    "open_interest",
+    "ask",
+    "bid",
+    "expiration_date",
+    "bid_date",
+    "volume",
+    "updated_at",
+    "gamma",
+    "delta",
+    "vega",
+    "volatility",
+]
+
 
 class JSONSerializableBaseModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, populate_by_name=True)
@@ -128,14 +149,33 @@ class OptionContract(JSONSerializableBaseModel):
     lastTradingDay: Optional[int] = None
     multiplier: Optional[float] = None
     settlementType: str  # AM, PM
-    isIndex: bool
+    isIndex: Optional[bool] = None
     percentChange: Optional[float] = None
     markChange: Optional[float] = None
     markPercentChange: Optional[float] = None
     model_config = ConfigDict(
         validate_assignment=False, use_enum_values=True, populate_by_name=True
     )
 
+    def to_dataframe_row(self) -> List[Any]:
+        result: List[Any] = [
+            self.strikePrice,
+            self.symbol.strip().replace(" ", ""),
+            self.lastPrice,
+            self.openInterest,
+            self.askPrice,
+            self.bidPrice,
+            self.expirationDate,
+            util.ts_to_date_string(self.quoteTimeInLong),
+            self.totalVolume,
+            util.ts_to_date_string(self.quoteTimeInLong),
+            self.gamma,
+            self.delta,
+            self.vega,
+            self.volatility,
+        ]
+        return result
+
 
 class Underlying(JSONSerializableBaseModel):
     ask: float
@@ -158,6 +198,18 @@ class Underlying(JSONSerializableBaseModel):
     totalVolume: Optional[int] = None
     tradeTime: Optional[int] = None
 
+    @property
+    def quote_time(self) -> datetime:
+        return util.ts_to_datetime(self.quoteTime)
+
+
+@dataclass
+class OptionChainDataFrames:
+    expiration: str
+    underlying_symbol: str
+    call_df: pd.DataFrame
+    put_df: pd.DataFrame
+
 
 class OptionChain(JSONSerializableBaseModel):
     symbol: str
@@ -178,3 +230,49 @@ class OptionChain(JSONSerializableBaseModel):
     callExpDateMap: Mapping[
         str, Mapping[str, List[OptionContract]]
     ]  # key: expiration:27 value:[strike: OptionContract]
+
+    def to_dataframe_pairs_by_expiration(self) -> List[OptionChainDataFrames]:
+        """
+        List of OptionChainDataFrames by expiration.
+        Each OptionChainDataFrames object contains call and put chain in dataframe format.
+        """
+        results: List[OptionChainDataFrames] = []
+        call_map = self.break_down_option_map(self.callExpDateMap)
+        put_map = self.break_down_option_map(self.putExpDateMap)
+        for expiration in call_map.keys():
+            call_df = call_map[expiration]
+            put_df = put_map[expiration]
+
+            cur_df_pair = OptionChainDataFrames(
+                expiration=expiration,
+                underlying_symbol=self.symbol,
+                call_df=call_df,
+                put_df=put_df,
+            )
+            results.append(cur_df_pair)
+
+        return results
+
+    def break_down_option_map(
+        self, optionExpMap: Mapping[str, Mapping[str, List[OptionContract]]]
+    ) -> Mapping[str, pd.DataFrame]:
+        result: MutableMapping[str, pd.DataFrame] = {}
+        for exp_date, strike_map in optionExpMap.items():
+            expiration = exp_date.split(":")[0]
+            if expiration not in result:
+                result[expiration] = {}
+
+            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()
+                    row.insert(0, self.underlying.mark)
+                    all_rows.append(row)
+
+                strike_df = pd.DataFrame(all_rows)
+                strike_df.columns = OptionChain_Headers
+            result[expiration] = strike_df
+
+        return result

cschwabpy/util.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 import pytz
+from typing import Optional
 
 eastern_tz: pytz.BaseTzInfo = pytz.timezone("US/Eastern")
 YMD_FMT = "%Y-%m-%d"
@@ -17,12 +18,24 @@ def now_unix_ts() -> float:
     return now().timestamp()
 
 
-def ts_to_datetime(ts: float, tz: pytz.BaseTzInfo = eastern_tz) -> datetime:
+def ts_to_datetime(
+    ts: Optional[float] = None, tz: pytz.BaseTzInfo = eastern_tz
+) -> Optional[datetime]:
+    if ts is None:
+        return None
     while ts > 1e10:
         ts = ts / 1000
     return datetime.fromtimestamp(ts, tz)
 
 
+def ts_to_date_string(
+    ts: Optional[float] = None, tz: pytz.BaseTzInfo = eastern_tz
+) -> Optional[str]:
+    if ts is None:
+        return None
+    return ts_to_datetime(ts, tz).strftime(YMD_FMT)
+
+
 def today_str(tz: pytz.BaseTzInfo = eastern_tz) -> str:  # type: ignore
     """Today in string. Returns Y-m-d."""  # noqa: DAR201
     return now(tz=tz).strftime(YMD_FMT)

tests/test_models.py

@@ -27,8 +27,11 @@ def test_option_chain_parsing() -> None:
     assert opt_chain_result is not None
     assert opt_chain_result.status == "SUCCESS"
 
-    # for key, value in opt_chain_result.callExpDateMap.items():
-    #     # Do something with key and value
-    #     print(key[:10])
-    #     print(value)#
-    #     print("----------------")
+    opt_df_pairs = opt_chain_result.to_dataframe_pairs_by_expiration()
+    assert opt_df_pairs is not None
+    for df in opt_df_pairs:
+        print(df.expiration)
+        print(f"call dataframe size: {df.call_df.shape}. expiration: {df.expiration}")
+        print(f"put dataframe size: {df.put_df.shape}. expiration: {df.expiration}")
+        print(df.call_df.head(5))
+        print(df.put_df.head(5))