feat(market_watch_init): use int64

index_col does not work with Int64 datatype, see:
pandas-dev/pandas#44835

Author: 5j9

test/test_tsetmc.py

@@ -303,26 +303,39 @@ def test_fmelli_instant():
 
 @patch_get_content('MarketWatchInit.aspx')
 def test_market_watch_init():
-    df = market_watch_init()['prices']
-    assert df.dtypes.to_dict() == {
+    dtypes = {
         'heven': 'int64', 'pf': 'int64', 'pc': 'int64'
         , 'pl': 'int64', 'tno': 'int64', 'tvol': 'int64'
         , 'tval': 'int64', 'pmin': 'int64', 'pmax': 'int64'
         , 'py': 'int64', 'eps': 'float64', 'bvol': 'int64'
         , 'visitcount': 'int64', 'flow': 'int64'
         , 'cs': 'int64', 'tmax': 'float64', 'tmin': 'float64'
-        , 'z': 'int64', 'yval': 'int64', 'pd1': 'Int64'
-        , 'po1': 'Int64', 'qd1': 'Int64', 'qo1': 'Int64'
-        , 'zd1': 'Int64', 'zo1': 'Int64', 'pd2': 'Int64'
-        , 'po2': 'Int64', 'qd2': 'Int64', 'qo2': 'Int64'
-        , 'zd2': 'Int64', 'zo2': 'Int64', 'pd3': 'Int64'
-        , 'po3': 'Int64', 'qd3': 'Int64', 'qo3': 'Int64'
-        , 'zd3': 'Int64', 'zo3': 'Int64'}
-    index = df.index
-    assert index.names == ['ins_code', 'isin', 'l18', 'l30']
-    assert index.dtype == 'O'
-    # in pandas 1.3 there will be no need for `.to_frame()`
-    assert index.to_frame().dtypes.to_list() == ['int64', 'O', 'O', 'O']
+        , 'z': 'int64', 'yval': 'int64'}
+
+    mwi = market_watch_init(join=False, market_state=False)
+    assert mwi['prices'].dtypes.to_dict() == dtypes
+    assert mwi['best_limits'].index.names == ['ins_code', 'row']
+    assert 'market_state' not in mwi
+
+    mwi = market_watch_init(market_state=True)
+    prices = mwi['prices']
+    assert 'market_state' in mwi
+    assert 'best_limits' in mwi
+    assert prices.dtypes.to_dict() == {
+        **dtypes, 'pd1': 'int64'
+        , 'po1': 'int64', 'qd1': 'int64', 'qo1': 'int64'
+        , 'zd1': 'int64', 'zo1': 'int64', 'pd2': 'int64'
+        , 'po2': 'int64', 'qd2': 'int64', 'qo2': 'int64'
+        , 'zd2': 'int64', 'zo2': 'int64', 'pd3': 'int64'
+        , 'po3': 'int64', 'qd3': 'int64', 'qo3': 'int64'
+        , 'zd3': 'int64', 'zo3': 'int64'}
+    assert prices.index.dtypes.to_dict() == {
+        'ins_code': 'int64', 'isin': 'O', 'l18': 'O', 'l30': 'O'}
+
+    mwi = market_watch_init(prices=False, market_state=False)
+    assert 'prices' not in mwi
+    assert mwi['best_limits'].index.dtypes.to_dict() == {
+        'ins_code': 'int64', 'row': 'int64'}
 
 
 @patch_get_content('ClosingPriceAll.aspx')

tsetmc/_tsetmc.py

@@ -125,6 +125,9 @@ class IntraDay(TypedDict, total=False):
     best_limits: DataFrame
 
 
+PRICE_INDEX_COLS = ['ins_code', 'isin', 'l18', 'l30']
+
+
 class Instrument:
 
     __slots__ = 'code', '_l18', '_l30', '_cisin'
@@ -516,7 +519,7 @@ def adjustments(self) -> DataFrame:
 
 
 def market_watch_init(
-    market_state=False, prices=True, best_limits=True, join=True
+    *, market_state=True, prices=True, best_limits=True, join=True
 ) -> MarketWatchInit:
     """Return the market status which are the info used in creating filters.
 
@@ -534,7 +537,7 @@ def market_watch_init(
             (it's the time of last transaction in HHMMSS format)
     """
     text = fa_norm_text('http://tsetmc.com/tsev2/data/MarketWatchInit.aspx?h=0&r=0')
-    _, index_data, states, price_rows, _ = text.split('@')
+    _, market_state_str, states, price_rows, _ = text.split('@')
     result = {}
     if prices:
         result['prices'] = price_df = read_csv(
@@ -544,25 +547,25 @@ def market_watch_init(
                 'ins_code', 'isin', 'l18', 'l30', 'heven', 'pf', 'pc', 'pl',
                 'tno', 'tvol', 'tval', 'pmin', 'pmax', 'py', 'eps', 'bvol',
                 'visitcount' , 'flow', 'cs', 'tmax', 'tmin', 'z', 'yval'),
-            index_col=['ins_code', 'isin', 'l18', 'l30'])
+            index_col=PRICE_INDEX_COLS)
     if best_limits:
         result['best_limits'] = best_limits_df = read_csv(
             StringIO(price_rows), lineterminator=';', names=(
                 'ins_code', 'row', 'zo', 'zd', 'pd', 'po', 'qd', 'qo'),
-            dtype='Int64')
+            dtype='int64', index_col=('ins_code', 'row'))
+        # best_limits_df.set_index(['ins_code', 'row'], inplace=True)
     if join and prices and best_limits:
         # merge multiple rows sharing the same `row` number into one row.
         # a fascinating solution from https://stackoverflow.com/a/53563551/2705757
         # noinspection PyUnboundLocalVariable
-        best_limits_df.set_index(['ins_code', 'row'], inplace=True)
         best_limits_df = best_limits_df.unstack(fill_value=0).sort_index(axis=1, level=1)
         best_limits_df.columns = [f'{c}{i}' for c, i in best_limits_df.columns]
         # noinspection PyUnboundLocalVariable
         joined = best_limits_df.join(price_df)
         # joined_df.index = to_numeric(joined_df.index, downcast='unsigned')
         result['prices'] = joined
     if market_state:
-        result['market_state'] = _parse_market_state(index_data)
+        result['market_state'] = _parse_market_state(market_state_str)
     return result