added get dates opt

samzorps · samzorps · commit b585255be3fb · 2023-04-07T19:02:32.000-04:00
diff --git a/thetadata/client.py b/thetadata/client.py
@@ -1,6 +1,7 @@
 """Module that contains Theta Client class."""
 import datetime
 import threading
+import time
 import traceback
 from decimal import Decimal
 from threading import Thread
@@ -21,13 +22,13 @@
     Header,
     TickBody,
     ListBody,
-    parse_list_REST, parse_hist_REST
+    parse_list_REST, parse_flexible_REST, parse_hist_REST, parse_hist_REST_stream, parse_hist_REST_stream_ijson,
 )
 from .terminal import check_download, launch_terminal
 
 _NOT_CONNECTED_MSG = "You must establish a connection first."
 _VERSION = '0.9.0'
-
+URL_BASE = "http://localhost:25510/"
 
 def _format_strike(strike: float) -> int:
     """Round USD to the nearest tenth of a cent, acceptable by the terminal."""
@@ -657,7 +658,6 @@ def get_hist_option_REST(
         date_range: DateRange,
         interval_size: int = 0,
         use_rth: bool = True,
-        progress_bar: bool = False,
     ) -> pd.DataFrame:
         """
          Get historical options data.
@@ -671,7 +671,6 @@ def get_hist_option_REST(
         :param interval_size:  The interval size in milliseconds. Applicable to most requests except ReqType.TRADE.
         :param use_rth:        If true, timestamps prior to 09:30 EST and after 16:00 EST will be ignored
                                   (only applicable to intervals requests).
-        :param progress_bar:   Print a progress bar displaying download progress.
 
         :return:               The requested data as a pandas DataFrame.
         :raises ResponseError: If the request failed.
@@ -683,13 +682,18 @@ def get_hist_option_REST(
         start_fmt = _format_date(date_range.start)
         end_fmt = _format_date(date_range.end)
         right_fmt = right.value
-
+        use_rth_fmt = str(use_rth).lower()
         url = f"http://localhost:25510/hist/option/{req_fmt}"
         querystring = {"root": root, "start_date": start_fmt, "end_date": end_fmt,
                        "strike": strike_fmt, "exp": exp_fmt, "right": right_fmt,
-                       "ivl": interval_size}
+                       "ivl": interval_size, "rth": use_rth_fmt}
+        t1 = time.time()
         response = requests.get(url, params=querystring)
-        df = parse_hist_REST(response, use_rth, progress_bar)
+        t2 = time.time()
+        df = parse_flexible_REST(response)
+        t3 = time.time()
+        print(f'time for request.get: {t2-t1}')
+        print(f'time for parse_flexible_REST(): {t3-t2}')
         return df
 
     def get_opt_at_time(
@@ -774,7 +778,7 @@ def get_opt_at_time_REST(
         querystring = {"root": root, "start_date": start_fmt, "end_date": end_fmt, "strike": strike_fmt,
                        "exp": exp_fmt, "right": right_fmt, "ivl": ms_of_day}
         response = requests.get(url, params=querystring)
-        df = parse_hist_REST(response)
+        df = parse_flexible_REST(response)
         return df
 
     def get_stk_at_time(
@@ -842,8 +846,7 @@ def get_stk_at_time_REST(
         querystring = {"root": root_fmt, "start_date": start_fmt,
                        "end_date": end_fmt, "ivl": ms_of_day}
         response = requests.get(url, params=querystring)
-        print(response.url)
-        df = parse_hist_REST(response)
+        df = parse_flexible_REST(response)
         return df
 
     def get_hist_stock(
@@ -894,7 +897,6 @@ def get_hist_stock_REST(
             date_range: DateRange,
             interval_size: int = 0,
             use_rth: bool = True,
-            progress_bar: bool = False,
     ) -> pd.DataFrame:
         """
          Get historical stock data.
@@ -904,7 +906,6 @@ def get_hist_stock_REST(
         :param date_range:     The dates to fetch.
         :param interval_size:  The interval size in milliseconds. Applicable only to OHLC & QUOTE requests.
         :param use_rth:         If true, timestamps prior to 09:30 EST and after 16:00 EST will be ignored.
-        :param progress_bar:   Print a progress bar displaying download progress.
 
         :return:               The requested data as a pandas DataFrame.
         :raises ResponseError: If the request failed.
@@ -914,12 +915,12 @@ def get_hist_stock_REST(
         req_fmt = req.name.lower()
         start_fmt = _format_date(date_range.start)
         end_fmt = _format_date(date_range.end)
-
+        use_rth_fmt = str(use_rth).lower()
         url = f"http://localhost:25510/hist/stock/{req_fmt}"
-        querystring = {"root": root, "start_date": start_fmt, "end_date": end_fmt,
-                       "ivl": interval_size}
-        response = requests.get(url, params=querystring)
-        df = parse_hist_REST(response, use_rth, progress_bar)
+        params = {"root": root, "start_date": start_fmt, "end_date": end_fmt,
+                      "ivl": interval_size, "rth": use_rth_fmt}
+        response = requests.get(url, params=params)
+        df = parse_flexible_REST(response)
         return df
 
     # LISTING DATA
@@ -953,11 +954,13 @@ def get_dates_stk_REST(self, root: str, req: StockReqType) -> pd.Series:
         :raises ResponseError: If the request failed.
         :raises NoData:        If there is no data available for the request.
         """
-        url = "http://localhost:25510/list/dates/stock/quote"
-        params = {'root': root, 'req': req}
+        root_fmt = root.lower()
+        req_fmt = req.name.lower()
+        url = f"http://localhost:25510/list/dates/stock/{req_fmt}"
+        params = {'root': root_fmt}
         response = requests.get(url, params=params)
-        df = parse_list_REST(response, dates=True)
-        return df
+        series = parse_list_REST(response, dates=True)
+        return series
 
     def get_dates_opt(
             self,
@@ -1015,7 +1018,6 @@ def get_dates_opt_REST(
         sec = SecType.OPTION.value.lower()
         url = f"http://localhost:25510/list/dates/{sec}/{req}"
         params = {'root': root, 'exp': exp_fmt, 'strike': strike_fmt, 'right': right}
-        # TODO: try using pd.read_json(url) to directly get dataframe from the URL
         response = requests.get(url, params=params)
         df = parse_list_REST(response, dates=True)
         return df
@@ -1044,6 +1046,33 @@ def get_dates_opt_bulk(
         body = ListBody.parse(out, header, self._recv(header.size), dates=True)
         return body.lst
 
+    def get_dates_opt_bulk_REST(
+            self,
+            req: OptionReqType,
+            root: str,
+            exp: date) -> pd.Series:
+        """
+        Get all dates of data available for a given options contract and request type.
+
+        :param req:            The request type.
+        :param root:           The root / underlying / ticker / symbol.
+        :param exp:            The expiration date. Must be after the start of `date_range`.
+        :param strike:         The strike price in USD.
+        :param right:          The right of an options.
+
+        :return:               All dates that Theta Data provides data for given a request.
+        :raises ResponseError: If the request failed.
+        :raises NoData:        If there is no data available for the request.
+        """
+        req = req.name.lower()
+        exp_fmt = _format_date(exp)
+        sec = SecType.OPTION.value.lower()
+        url = f"http://localhost:25510/list/dates/{sec}/{req}"
+        params = {'root': root, 'exp': exp_fmt}
+        response = requests.get(url, params=params)
+        df = parse_list_REST(response, dates=True)
+        return df
+
     def get_expirations(self, root: str) -> pd.Series:
         """
         Get all options expirations for a provided underlying root.
@@ -1168,14 +1197,10 @@ def get_roots_REST(self, sec: SecType) -> pd.Series:
         :raises ResponseError: If the request failed.
         :raises NoData:        If there is no data available for the request.
         """
-        assert self._server is not None, _NOT_CONNECTED_MSG
         url = "http://localhost:25510/list/roots"
-        headers = {"Content-Type": "application/json"}
         params = {'sec': sec.value}
-        #make call
-        raw_json = requests.get(url, params=params).text
-        df = pd.read_json(raw_json, typ="series")
-        df = pd.Series(df['response'])
+        response = requests.get(url, params=params)
+        df = parse_list_REST(response)
         return df
 
     # LIVE DATA
diff --git a/thetadata/parsing.py b/thetadata/parsing.py
@@ -2,6 +2,12 @@
 from __future__ import annotations
 
 import json
+import urllib
+from urllib.parse import urlencode
+from urllib.request import urlopen
+
+import ijson
+import time
 from typing import Optional
 
 import requests
@@ -120,7 +126,7 @@ def _check_body_errors(header: Header, body_data: bytes):
             raise ResponseError(msg)
 
 
-def _check_header_errors_REST(response: requests.Response, header: dict):
+def _check_header_errors_REST(header: dict):
     """Check for errors from the Terminal.
 
     :raises NoData: if the server does not contain data for the request.
@@ -272,16 +278,37 @@ def _post_process(cls, df: DataFrame) -> None:
             )
 
 
-def parse_hist_REST(response: requests.Response, use_rth=False, progress_bar=False
-                    ) -> pd.DataFrame:
+def parse_flexible_REST(response: requests.Response) -> pd.DataFrame:
+    """
+    Flexible parsing function that uses a python dictionary as an intermediary
+    between json string and pandas dataframe.
+    """
     response_dict = response.json()
-    _check_header_errors_REST(response, response_dict["header"])
-
-    # Get column names from header "format" field and map them to correct type
-    cols = [DataType.from_string(name=c) for c in response_dict['header']['format']]
-    # Get the data rows from the "response" field
+    _check_header_errors_REST(response_dict["header"])
+    cols = [DataType.from_string(name=col) for col in response_dict['header']['format']]
     rows = response_dict['response']
-    # Create a pandas DataFrame
+    df = pd.DataFrame(rows, columns=cols)
+    if DataType.DATE in df.columns:
+        df[DataType.DATE] = pd.to_datetime(
+            df[DataType.DATE], format="%Y%m%d"
+        )
+    try:
+        return df
+    except Exception as e:
+        raise ResponseParseError(
+            f"Failed to parse header for request: {response.url}. Please send this error to support."
+        ) from e
+
+
+def parse_hist_REST(response: requests.Response) -> pd.DataFrame:
+    resp_split = response.text.split('"response": ')
+    to_lstrip = '"header": \t\n'
+    to_rstrip = ", \t\n"
+    header_str = resp_split[0][1:].lstrip(to_lstrip).rstrip(to_rstrip)
+    header = json.loads(header_str)
+    _check_header_errors_REST(header)
+    cols = [DataType.from_string(name=col) for col in header['format']]
+    rows = pd.read_json(resp_split[1][:-1], orient="table")
     df = pd.DataFrame(rows, columns=cols)
     if DataType.DATE in df.columns:
         df[DataType.DATE] = pd.to_datetime(
@@ -296,6 +323,60 @@ def parse_hist_REST(response: requests.Response, use_rth=False, progress_bar=Fal
         ) from e
 
 
+def parse_hist_REST_stream_ijson(url, params) -> pd.DataFrame:
+    url = url + '?' + urlencode(params)
+    f = urlopen(url)
+    header = {}
+    row = []
+    header_format = []
+    loc = 0
+    for prefix, event, value in ijson.parse(f, use_float=True):
+        if prefix == "response.item.item":
+            row.append(value)
+
+        elif prefix == "response.item" and event == "end_array":
+            df.loc[loc] = row
+            loc += 1
+            row = []
+
+        elif prefix == "header.format.item":
+            header_format.append(value)
+
+        elif prefix[:6] == "header" and len(prefix) > 6:
+            header[prefix[7:]] = value
+
+        elif event == "map_key" and value == "response":
+            header["format"] = header_format
+            _check_header_errors_REST(header)
+            cols = [DataType.from_string(name=col) for col in header['format']]
+            df = pd.DataFrame(columns=cols)
+            
+    if DataType.DATE in df.columns:
+        df[DataType.DATE] = pd.to_datetime(
+            df[DataType.DATE], format="%Y%m%d"
+        )
+    try:
+        return df
+    except Exception as e:
+        raise ResponseParseError(
+            f"Failed to parse header for request: {url}. Please send this error to support."
+        ) from e
+
+
+def parse_hist_REST_stream(url, params) -> pd.DataFrame:
+    header = {}
+    row = []
+    header_format = []
+    loc = 0
+    s = requests.Session()
+    with requests.get(url, params=params, stream=True) as resp:
+        line_num = 0
+        for line in resp.iter_lines():
+            print(line)
+            line_num += 1
+            if line_num > 10: break
+
+
 class ListBody:
     """Represents the body returned on every Terminal call that have one DataType."""
 
@@ -342,9 +423,7 @@ def _parse(
         return cls(lst=lst)
 
 
-def parse_list_REST(
-        response: requests.Response, dates: bool = False
-) -> pd.Series:
+def parse_list_REST(response: requests.Response, dates: bool = False) -> pd.Series:
     """Parse binary body data into an object.
 
     :param response: the requests.Response object
@@ -353,15 +432,13 @@ def parse_list_REST(
     """
     df = pd.read_json(response.text, typ="series")
     header = df['header']
-    url = response.history[0].url if response.history else response.url
-    _check_header_errors_REST(response, header)
-
+    _check_header_errors_REST(header)
     try:
-        df = pd.Series(df['response'])
+        df = pd.Series(df['response'], copy=False)
         if dates:
             df = pd.to_datetime(df, format="%Y%m%d")
         return df
     except Exception as e:
         raise ResponseParseError(
-            f"Failed to parse request: {url}. Please send this error to support."
+            f"Failed to parse request: {response.url}. Please send this error to support."
         ) from e
