1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-18 01:44:34 +08:00

Pass mypy

This commit is contained in:
Huoran Li
2022-07-26 14:56:24 +08:00
parent 83d8f00a19
commit ccc3f96ed3
7 changed files with 39 additions and 34 deletions

View File

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

View File

@@ -15,8 +15,8 @@ from qlib.data.dataset import DatasetH
class LRUCache: class LRUCache:
def __init__(self, pool_size: int = 200): def __init__(self, pool_size: int = 200):
self.pool_size = pool_size self.pool_size = pool_size
self.contents = dict() self.contents: dict = {}
self.keys = collections.deque() self.keys: collections.deque = collections.deque()
def put(self, key, item): def put(self, key, item):
if self.has(key): if self.has(key):
@@ -52,7 +52,7 @@ class DataWrapper:
self.feature_cache = LRUCache() self.feature_cache = LRUCache()
self.backtest_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) start_time, end_time = date.replace(hour=0, minute=0, second=0), date.replace(hour=23, minute=59, second=59)
if backtest: if backtest:

View File

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

View File

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

View File

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

View File

@@ -84,7 +84,7 @@ class DataQueue(Generic[T]):
self.activate() self.activate()
return self return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None: def __exit__(self, exc_type, exc_val, exc_tb):
self.cleanup() self.cleanup()
def cleanup(self) -> None: def cleanup(self) -> None:

View File

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