quantification 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

quantification/core/asset/stock.py

@@ -0,0 +1,268 @@
+from abc import ABCMeta, abstractmethod
+from typing import overload, Type
+from datetime import date
+from itertools import chain
+
+import akshare as ak
+
+from .base_asset import BaseAsset
+
+from ..cache import cache_query
+
+
+class BaseStockExchangeMeta(ABCMeta):
+    def __repr__(self: "BaseStockExchange"):
+        return self.name()
+
+    __str__ = __repr__
+
+
+class BaseStockExchange(metaclass=BaseStockExchangeMeta):
+    @classmethod
+    @abstractmethod
+    def name(cls) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def code(cls) -> str:
+        raise NotImplementedError
+
+
+# 上海证券交易所
+class SSE(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "上海证券交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "SH"
+
+
+# 深圳证券交易所
+class SZSE(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "深圳证券交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "SZ"
+
+
+# 北京证券交易所
+class BSE(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "北京证券交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "BJ"
+
+
+# 香港交易所
+class HKEX(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "香港交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "HK"
+
+
+# 纽约证券交易所
+class NYSE(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "纽约证券交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "N"
+
+
+# 纳斯达克
+class NASDAQ(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "纳斯达克"
+
+    @classmethod
+    def code(cls) -> str:
+        return "O"
+
+
+# 伦敦证券交易所
+class LSE(BaseStockExchange):
+    @classmethod
+    def name(cls) -> str:
+        return "伦敦证券交易所"
+
+    @classmethod
+    def code(cls) -> str:
+        return "L"
+
+
+class StockExchange:
+    SSE = SSE
+    SZSE = SZSE
+    BSE = BSE
+    HKEX = HKEX
+    NYSE = NYSE
+    NASDAQ = NASDAQ
+    LSE = LSE
+
+
+def predict_exchange(symbol: str) -> type[BaseStockExchange]:
+    if "." not in symbol:
+        api = cache_query(update=False)(ak.stock_zh_a_spot)
+        code_list = api()["代码"].tolist()
+        for code in code_list:
+            if code[2:] != symbol: continue
+
+            match code[:2]:
+                case "sh":
+                    return StockExchange.SSE
+                case "sz":
+                    return StockExchange.SZSE
+                case "bj":
+                    return StockExchange.BSE
+
+        raise ValueError(f"无法根据{symbol}判断交易所")
+
+    _, exchange = symbol.split(".")
+    match exchange.lower():
+        case "sh":
+            return StockExchange.SSE
+        case "sz":
+            return StockExchange.SZSE
+        case "bj":
+            return StockExchange.BSE
+        case _:
+            raise ValueError(f"无法根据{symbol}判断交易所")
+
+
+stock_family: dict[str:type["Stock"]] = {}
+SharePosition = dict[date, int]
+
+
+class Stock(BaseAsset):
+    symbol: str = None
+    exchange: type[BaseStockExchange] = None
+
+    @classmethod
+    def type(cls, *args, **kwargs):
+        return "Stock"
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return f"股票{cls.symbol}.{cls.exchange.code()}" if cls.symbol else "股票"
+
+    @property
+    def amount(self, *args, **kwargs):
+        return sum(self.share_position.values())
+
+    @property
+    def extra(self, *args, **kwargs):
+        return {"position": {k.isoformat(): int(v) for k, v in self.share_position.items()}}
+
+    @property
+    def is_empty(self):
+        return self.amount == 0
+
+    @property
+    def copy(self):
+        return Stock[self.symbol](self.share_position)
+
+    def available(self, day: date):
+        available_share = 0
+
+        for brought_day, share in self.share_position.items():
+            if brought_day < day:
+                available_share += share
+
+        return Stock[self.symbol](available_share)
+
+    def __add__(self, other: "Stock"):
+        assert self == other, f"只有同种股票可以相加减, {self.__class__.__name__} != {other.__class__.__name__}"
+
+        return Stock[self.symbol](self.add_shares(other.share_position))
+
+    def __sub__(self, other: "Stock"):
+        assert self == other, f"只有同种股票可以相加减, {self.__class__.__name__} != {other.__class__.__name__}"
+
+        day = next(iter(other.share_position.keys()))
+        available = self.available(day).amount
+        assert available >= other.amount, f"可出售股份不足, {available} < {other.amount}"
+
+        return Stock[self.symbol](self.sub_shares(other.amount))
+
+    def __eq__(self, other):
+        return isinstance(other, Stock) and self.symbol == other.symbol
+
+    def __repr__(self):
+        return f"{self.name()}({self.amount}股)"
+
+    __str__ = __repr__
+
+    def __init__(self, share: int | SharePosition):
+        assert self.symbol is not None, "未指定股票代码, 无法实例化, 是否应该使用Stock[symbol](...) ?"
+
+        if isinstance(share, int):
+            assert share >= 0, f"股票份数不能为负, 实际为{share=}"
+            self.share_position: SharePosition = {self.env.date: share}
+        else:
+            assert all([v >= 0 for v in share.values()]), f"股票份数不能为负, 实际为{share.values()=}"
+            self.share_position = share
+
+    def add_shares(self, share_position: SharePosition) -> SharePosition:
+        return {
+            k: self.share_position.get(k, 0) + share_position.get(k, 0)
+            for k in chain(self.share_position, share_position)
+        }
+
+    @overload
+    def __class_getitem__(cls, symbol: str):
+        ...
+
+    @overload
+    def __class_getitem__(cls, symbol_and_exchange: tuple[str, Type[BaseStockExchange]]):
+        ...
+
+    def __class_getitem__(cls, arg: str | tuple[str, Type[BaseStockExchange]]):
+        if isinstance(arg, str):
+            symbol = arg
+            exchange = predict_exchange(arg)
+        else:
+            symbol, exchange = arg
+
+        if not stock_family.get(f"{symbol}.{exchange}"):
+            stock_family[symbol] = type(
+                f"Stock{symbol}",
+                (Stock,),
+                {"symbol": symbol, "exchange": exchange},
+            )
+
+        return stock_family[symbol]
+
+    def sub_shares(self, share: int) -> SharePosition:
+        remaining_share = share
+        new_share_position = {}
+        sorted_dates = sorted(self.share_position.keys())
+        for sorted_date in sorted_dates:
+            current_share = self.share_position[sorted_date]
+            if remaining_share <= 0:
+                new_share_position[sorted_date] = current_share
+                continue
+            if current_share <= remaining_share:
+                remaining_share -= current_share
+            else:
+                new_share_position[sorted_date] = current_share - remaining_share
+                remaining_share = 0
+        return new_share_position
+
+
+__all__ = ["Stock", "StockExchange"]

quantification/core/cache.py

@@ -0,0 +1,93 @@
+import time
+import pickle
+import hashlib
+import inspect
+from pathlib import Path
+from functools import partial
+
+from .configure import config
+
+CACHE_DIR = Path(config.cache_dir)
+CACHE_QUERY = CACHE_DIR / 'query'
+
+CACHE_QUERY.mkdir(parents=True, exist_ok=True)
+
+
+def flush(cache_file, func, *args, **kwargs):
+    result = func(*args, **kwargs)
+    with open(cache_file, 'wb') as file:
+        pickle.dump(result, file)
+    return result
+
+
+def get_or_flush(cache_file, func, *args, **kwargs):
+    try:
+        with open(cache_file, 'rb') as f:
+            return pickle.load(f)
+    except (pickle.UnpicklingError, EOFError, FileNotFoundError):
+        return flush(cache_file, func, *args, **kwargs)
+
+
+def cache_query(update=None, expire_seconds=86400):
+    """缓存装饰器, 通过参数控制缓存更新逻辑
+
+    Args:
+        update (bool|None):
+            - True: 强制更新缓存
+            - False: 只要缓存存在就使用（忽略有效期）
+            - None: 根据 `expire_seconds` 判断是否更新
+        expire_seconds (int): 缓存有效期（秒）, 仅在 update=None 时生效
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+
+            # 绑定参数生成唯一缓存键
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            args_dict = bound_args.arguments
+            sorted_args = sorted(args_dict.items(), key=lambda x: x[0])
+            name = func.__name__ if not isinstance(func, partial) else func.func.__name__
+            key_data = (name, sorted_args)
+
+            # 计算哈希作为文件名
+            hash_key = hashlib.sha256(pickle.dumps(key_data)).hexdigest()
+            cache_file = CACHE_QUERY / f"{hash_key}.pkl"
+            cache_exists = cache_file.exists()
+
+            match update:
+                case True:
+                    return flush(cache_file, func, *args, **kwargs)
+
+                case False:
+                    if not cache_exists:
+                        return flush(cache_file, func, *args, **kwargs)
+
+                    return get_or_flush(cache_file, func, *args, **kwargs)
+
+                case None:
+                    if not cache_exists:
+                        return flush(cache_file, func, *args, **kwargs)
+
+                    if expire_seconds is not None:
+                        current_time = time.time()
+                        mtime = cache_file.stat().st_mtime
+                        cache_valid = (current_time - mtime) <= expire_seconds
+                    else:
+                        cache_valid = True  # 无有效期要求
+
+                    if not cache_valid:
+                        return flush(cache_file, func, *args, **kwargs)
+
+                    return get_or_flush(cache_file, func, *args, **kwargs)
+
+                case _:
+                    raise ValueError("Invalid update parameter value")
+
+        return wrapper
+
+    return decorator
+
+
+__all__ = ['cache_query']

quantification/core/configure.py

@@ -0,0 +1,15 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Config(BaseSettings):
+    model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8', extra="allow")
+
+    log_level: str = Field(default="INFO", description="日志等级")
+    cache_dir: str = Field(default="./cache", description="缓存文件夹位置")
+    adjust: str = Field(description='复权方式')
+
+
+config = Config()
+
+__all__ = ["Config", "config"]

quantification/core/data/__init__.py

@@ -0,0 +1,5 @@
+from .base_api import BaseAPI, BaseCombinedAPI
+from .base_delegate import BaseDelegate
+
+from .field import Field
+from .panel import DataPanel

quantification/core/data/base_api.py

@@ -0,0 +1,109 @@
+from typing import TypeVar, Generic
+from datetime import date
+from itertools import product
+
+from pydantic import BaseModel
+
+from .field import Field
+from .panel import DataPanel
+from .base_delegate import BaseDelegate
+
+from ..logger import logger
+from ..configure import config, Config
+
+SettingType = TypeVar('SettingType', bound=BaseModel)
+
+
+class BaseAPI(Generic[SettingType]):
+    setting_class: type[BaseModel]
+    delegate_classes: list[type[BaseDelegate[SettingType]]]
+
+    def __init_subclass__(cls, **kwargs):
+        assert hasattr(cls, "setting_class"), \
+            f"{cls.__name__}必须实现类属性setting_class"
+
+        assert issubclass(cls.setting_class, BaseModel), \
+            f"{cls.__name__}类属性setting_class必须为BaseModel子类, 实际为{cls.setting_class}"
+
+        assert hasattr(cls, "delegate_classes"), \
+            f"{cls.__name__}必须实现类属性delegate_classes"
+
+        assert isinstance(cls.delegate_classes, list), \
+            f"{cls.__name__}类属性delegate_classes必须为列表, 实际为{type(cls.delegate_classes)}"
+
+        for delegate_class in cls.delegate_classes:
+            assert issubclass(delegate_class, BaseDelegate), \
+                f"{cls.__name__}类属性delegate_classes的元素必须为BaseDelegate子类, 实际为{delegate_class}"
+
+    def __init__(self):
+        self.config = config
+        self.setting: SettingType = self.setting_class.model_validate(config.model_dump())
+        self.delegates = self.initialize_delegates()
+
+    def initialize_delegates(self):
+        return [i(self.config, self.setting) for i in self.delegate_classes]
+
+    def query(self, start_date: date, end_date: date, fields: list[Field], **kwargs):
+        fields_delegation: dict[BaseDelegate[SettingType], list[Field]] = {}
+
+        picked = []
+        for delegate, field in product(self.delegates, fields):
+            if field in picked:
+                continue
+
+            if not delegate.has_field(field, **kwargs):
+                continue
+
+            if not fields_delegation.get(delegate):
+                fields_delegation[delegate] = []
+
+            fields_delegation[delegate].append(field)
+            picked.append(field)
+
+        omitted = set(fields) - set(picked)
+        if omitted:
+            logger.warning(f"无法查询的字段: {omitted}")
+
+        res = DataPanel()
+
+        for delegate, fields in fields_delegation.items():
+            panel = delegate.execute(start_date, end_date, fields, **kwargs)
+            res <<= panel
+
+        return res
+
+    def __repr__(self):
+        return f"<API {self.__class__.__name__}>"
+
+    __str__ = __repr__
+
+
+class BaseCombinedAPI(BaseAPI):
+    setting_class = Config
+    delegate_classes = []
+
+    api_classes: list[type[BaseAPI]]
+
+    def __init_subclass__(cls, **kwargs):
+        assert hasattr(cls, "api_classes"), \
+            f"{cls.__name__}必须实现类属性api_classes"
+
+        for api_class in cls.api_classes:
+            assert issubclass(api_class, BaseAPI), \
+                f"{cls.__name__}类属性delegate_classes元素必须为BaseAPI子类, 实际为{api_class}"
+
+    def initialize_delegates(self):
+        delegates = []
+
+        for api_class in self.api_classes:
+            delegates += api_class().delegates
+
+        return delegates
+
+    def __repr__(self):
+        return f"<CombinedAPI {'|'.join(map(lambda x: x.__name__, self.api_classes))}>"
+
+    __str__ = __repr__
+
+
+__all__ = ["BaseAPI", "BaseCombinedAPI"]

quantification/core/data/base_delegate.py

@@ -0,0 +1,73 @@
+from abc import ABC, abstractmethod
+from typing import TypeVar, Generic, Callable
+from datetime import date
+
+import pandas as pd
+from pydantic import BaseModel
+
+from .field import Field
+from .panel import DataPanel
+
+from ..configure import Config
+
+SettingType = TypeVar('SettingType', bound=BaseModel)
+
+
+class BaseDelegate(ABC, Generic[SettingType]):
+    pair: list[tuple[Field, str]]
+    field2str: dict[Field, str]
+    str2field: dict[str, Field]
+
+    DATE_FIELD = "日期"
+
+    def __init_subclass__(cls, **kwargs):
+        if ABC in cls.__bases__: return
+
+        assert hasattr(cls, 'pair'), f'{cls.__name__}必须实现类属性pair'
+
+        cls.field2str = {k: v for k, v in cls.pair}
+        cls.str2field = {v: k for k, v in cls.pair}
+
+    def __init__(self, config: Config, setting: SettingType) -> None:
+        self.config = config
+        self.setting = setting
+
+    def rename_columns(self, data: pd.DataFrame, date_field: str) -> pd.DataFrame:
+        data = data.rename(columns=self.str2field)
+        data = data.rename(columns={date_field: self.DATE_FIELD})
+
+        return data
+
+    def use_date_index(self, data: pd.DataFrame, formatter: Callable = None) -> pd.DataFrame:
+        if formatter:
+            data[self.DATE_FIELD] = data[self.DATE_FIELD].apply(formatter)
+
+        data[self.DATE_FIELD] = pd.to_datetime(data[self.DATE_FIELD])
+        data = data.set_index(self.DATE_FIELD)
+        data = data.sort_index()
+
+        return data
+
+    def execute(self, start_date: date, end_date: date, fields: list[Field], **kwargs) -> DataPanel:
+        data = self.query(start_date, end_date, fields, **kwargs)
+        mask = self.mask(data, start_date, end_date, fields, **kwargs)
+        data = data[fields]
+        mask = mask[fields]
+        data = data[start_date:end_date]
+        mask = mask[start_date:end_date]
+        return DataPanel(data, mask)
+
+    @abstractmethod
+    def has_field(self, field: Field, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def query(self, start_date: date, end_date: date, fields: list[Field], **kwargs) -> pd.DataFrame:
+        raise NotImplementedError
+
+    @abstractmethod
+    def mask(self, data: pd.DataFrame, start_date: date, end_date: date, fields: list[Field], **kwargs) -> pd.DataFrame:
+        raise NotImplementedError
+
+
+__all__ = ["BaseDelegate"]