Pass mypy

Author: lihuoran

qlib/rl/from_neutrader/config.py

@@ -9,7 +9,7 @@
 @dataclass
 class ExchangeConfig:
     limit_threshold: Union[float, Tuple[str, str]]
-    deal_price: Union[str, Tuple[str, str]]
+    deal_price: Union[str, Tuple[str]]
     volume_threshold: dict
     open_cost: float = 0.0005
     close_cost: float = 0.0015

qlib/rl/from_neutrader/feature.py

@@ -15,8 +15,8 @@
 class LRUCache:
     def __init__(self, pool_size: int = 200):
         self.pool_size = pool_size
-        self.contents = dict()
-        self.keys = collections.deque()
+        self.contents: dict = {}
+        self.keys: collections.deque = collections.deque()
 
     def put(self, key, item):
         if self.has(key):
@@ -52,7 +52,7 @@ def __init__(
         self.feature_cache = LRUCache()
         self.backtest_cache = LRUCache()
 
-    def get(self, stock_id: str, date: pd.Timestamp, backtest: bool = False):
+    def get(self, stock_id: str, date: pd.Timestamp, backtest: bool = False) -> pd.DataFrame:
         start_time, end_time = date.replace(hour=0, minute=0, second=0), date.replace(hour=23, minute=59, second=59)
 
         if backtest:

qlib/rl/order_execution/interpreter.py

@@ -165,13 +165,11 @@ def interpret(self, state: SAOEState) -> CurrentStateObs:
         assert self.env is not None
         assert self.env.status["cur_step"] <= self.max_step
         obs = CurrentStateObs(
-            **{
-                "acquiring": state.order.direction == state.order.BUY,
-                "cur_step": self.env.status["cur_step"],
-                "num_step": self.max_step,
-                "target": state.order.amount,
-                "position": state.position,
-            }
+            acquiring=state.order.direction == state.order.BUY,
+            cur_step=self.env.status["cur_step"],
+            num_step=self.max_step,
+            target=state.order.amount,
+            position=state.position,
         )
         return obs
 

qlib/rl/order_execution/simulator_qlib.py

@@ -4,7 +4,7 @@
 """Placeholder for qlib-based simulator."""
 from __future__ import annotations
 
-from typing import Callable, Generator, List, Optional, Tuple, cast
+from typing import Any, Callable, Generator, List, Optional, Tuple, cast
 
 import numpy as np
 import pandas as pd
@@ -36,7 +36,7 @@ def __init__(self) -> None:
         self.execute_order: Optional[Order] = None
         self.execute_result: List[Tuple[Order, float, float, float]] = []
 
-    def generate_trade_decision(self, execute_result: list = None) -> BaseTradeDecision:
+    def generate_trade_decision(self, execute_result: list = None) -> Generator[Any, Any, BaseTradeDecision]:
         exec_vol = yield self
 
         oh = self.trade_exchange.get_order_helper()
@@ -52,7 +52,7 @@ def alter_outer_trade_decision(self, outer_trade_decision: BaseTradeDecision) ->
     def post_exe_step(self, execute_result: list) -> None:
         self.execute_result = execute_result
 
-    def reset(self, outer_trade_decision: TradeDecisionWO = None, **kwargs) -> None:
+    def reset(self, outer_trade_decision: TradeDecisionWO = None, **kwargs: Any) -> None:
         super().reset(outer_trade_decision=outer_trade_decision, **kwargs)
         if outer_trade_decision is not None:
             order_list = outer_trade_decision.order_list
@@ -83,7 +83,7 @@ def generate_trade_decision(self, execute_result: list = None) -> TradeDecisionW
             oh.create(
                 code=self._instrument,
                 amount=self._order.amount,
-                direction=Order.parse_dir(self._order.direction),
+                direction=self._order.direction,
             ),
         ]
         return TradeDecisionWO(order_list, self, self._trade_range)
@@ -102,7 +102,7 @@ def __init__(self, order: Order, tick_index: pd.DatetimeIndex, twap_price: float
         # NOTE: can empty dataframe contain index?
         self.history_exec = pd.DataFrame(columns=metric_keys).set_index("datetime")
         self.history_steps = pd.DataFrame(columns=metric_keys).set_index("datetime")
-        self.metrics = None
+        self.metrics: Optional[SAOEMetrics] = None
 
     def update(
         self,
@@ -116,6 +116,8 @@ def update(
         exec_vol = np.array([e[0].deal_amount for e in execute_result])
         num_step = len(execute_result)
 
+        assert execute_order is not None
+
         if num_step == 0:
             market_volume = np.array([])
             market_price = np.array([])
@@ -251,7 +253,7 @@ def __init__(
         exchange_config: ExchangeConfig,
     ) -> None:
         super().__init__(
-            initial=None,  # TODO
+            initial=order,  # TODO: confirm this logic
         )
 
         assert order.start_time.date() == order.end_time.date()
@@ -330,6 +332,8 @@ def reset(self, order: Order) -> None:
         )
 
     def _iter_strategy(self, action: float = None) -> DecomposedStrategy:
+        assert self._collect_data_loop is not None
+
         strategy = next(self._collect_data_loop) if action is None else self._collect_data_loop.send(action)
         while not isinstance(strategy, DecomposedStrategy):
             strategy = next(self._collect_data_loop) if action is None else self._collect_data_loop.send(action)
@@ -344,6 +348,7 @@ def step(self, action: float) -> None:
         except StopIteration:
             self._done = True
 
+        assert self._executor is not None
         _, all_indicators = get_portfolio_and_indicator(self._executor)
 
         self._maintainer.update(

qlib/rl/reward.py

@@ -31,6 +31,7 @@ def reward(self, simulator_state: SimulatorState) -> float:
         raise NotImplementedError("Implement reward calculation recipe in `reward()`.")
 
     def log(self, name: str, value: Any) -> None:
+        assert self.env is not None
         self.env.logger.add_scalar(name, value)
 
 

qlib/rl/utils/data_queue.py

@@ -84,7 +84,7 @@ def __enter__(self) -> DataQueue:
         self.activate()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(self, exc_type, exc_val, exc_tb):
         self.cleanup()
 
     def cleanup(self) -> None:

qlib/rl/utils/finite_env.py

@@ -57,9 +57,10 @@ def fill_invalid(obj: int | float | bool | np.ndarray | dict | list | tuple) ->
 
 def is_invalid(arr: int | float | bool | np.ndarray | dict | list | tuple) -> bool:
     if hasattr(arr, "dtype"):
-        if np.issubdtype(arr.dtype, np.floating):
+        dtype = getattr(arr, "dtype")
+        if np.issubdtype(dtype, np.floating):
             return np.isnan(arr).all()
-        return (np.iinfo(arr.dtype).max == arr).all()
+        return (np.iinfo(dtype).max == arr).all()
     if isinstance(arr, dict):
         return all(is_invalid(o) for o in arr.values())
     if isinstance(arr, (list, tuple)):
@@ -209,7 +210,7 @@ def collector_guard(self) -> Generator[FiniteVectorEnv, None, None]:
 
     def reset(
         self,
-        id: int | List[int] | np.ndarray = None,
+        id: int | List[int] | np.ndarray | None = None,
     ) -> np.ndarray:
         assert not self._zombie
 
@@ -222,23 +223,23 @@ def reset(
                 RuntimeWarning,
             )
 
-        id = self._wrap_id(id)
+        wrapped_id = self._wrap_id(id)
         self._reset_alive_envs()
 
         # ask super to reset alive envs and remap to current index
-        request_id = list(filter(lambda i: i in self._alive_env_ids, id))
-        obs = [None] * len(id)
-        id2idx = {i: k for k, i in enumerate(id)}
+        request_id = [i for i in wrapped_id if i in self._alive_env_ids]
+        obs = [None] * len(wrapped_id)
+        id2idx = {i: k for k, i in enumerate(wrapped_id)}
         if request_id:
             for i, o in zip(request_id, super().reset(request_id)):
                 obs[id2idx[i]] = self._postproc_env_obs(o)
 
-        for i, o in zip(id, obs):
+        for i, o in zip(wrapped_id, obs):
             if o is None and i in self._alive_env_ids:
                 self._alive_env_ids.remove(i)
 
         # logging
-        for i, o in zip(id, obs):
+        for i, o in zip(wrapped_id, obs):
             if i in self._alive_env_ids:
                 for logger in self._logger:
                     logger.on_env_reset(i, obs)
@@ -251,7 +252,7 @@ def reset(
                 obs[i] = self._get_default_obs()
 
         if not self._alive_env_ids:
-            # comment this line so that the env becomes indisposable
+            # comment this line so that the env becomes indispensable
             # self.reset()
             self._zombie = True
             raise StopIteration
@@ -261,13 +262,13 @@ def reset(
     def step(
         self,
         action: np.ndarray,
-        id: int | List[int] | np.ndarray = None,
+        id: int | List[int] | np.ndarray | None = None,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         assert not self._zombie
-        id = self._wrap_id(id)
-        id2idx = {i: k for k, i in enumerate(id)}
-        request_id = list(filter(lambda i: i in self._alive_env_ids, id))
-        result = [[None, None, False, None] for _ in range(len(id))]
+        wrapped_id = self._wrap_id(id)
+        id2idx = {i: k for k, i in enumerate(wrapped_id)}
+        request_id = list(filter(lambda i: i in self._alive_env_ids, wrapped_id))
+        result = [[None, None, False, None] for _ in range(len(wrapped_id))]
 
         # ask super to step alive envs and remap to current index
         if request_id:
@@ -277,7 +278,7 @@ def step(
                 result[id2idx[i]][0] = self._postproc_env_obs(result[id2idx[i]][0])
 
         # logging
-        for i, r in zip(id, result):
+        for i, r in zip(wrapped_id, result):
             if i in self._alive_env_ids:
                 for logger in self._logger:
                     logger.on_env_step(i, *r)