lidb 1.1.10__tar.gz

lidb-1.1.10/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: lidb
+Version: 1.1.10
+Summary: Light database for quantor
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: dynaconf>=3.2.11
+Requires-Dist: polars>=1.31.0
+Requires-Dist: sqlparse>=0.5.3
+Requires-Dist: logair>=1.0.1
+Requires-Dist: clickhouse-df>=0.1.5
+Requires-Dist: connectorx>=0.4.3
+Requires-Dist: pymysql>=1.1.2
+Requires-Dist: xcals>=0.0.4
+Requires-Dist: ygo>=1.2.6
+Requires-Dist: lark>=1.3.1
+Requires-Dist: numpy>=2.3.1

lidb-1.1.10/lidb/__init__.py

@@ -0,0 +1,30 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+
+from .init import (
+    NAME,
+    DB_PATH,
+    CONFIG_PATH,
+    get_settings,
+)
+
+from .database import (
+    sql,
+    put,
+    has,
+    tb_path,
+    read_mysql,
+    write_mysql,
+    execute_mysql,
+    read_ck,
+    scan,
+)
+
+from .table import Table, TableMode
+from .dataset import Dataset, load_ds
+from .qdf import to_lazy, Expr
+from .svc import DataService, D
+
+from .parse import parse_hive_partition_structure
+
+__version__ = "1.1.10"

lidb-1.1.10/lidb/database.py

@@ -0,0 +1,234 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Copyright (c) 2025 ZhangYundi
+Licensed under the MIT License.
+Created on 2024/7/1 09:44
+Email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+import re
+from pathlib import Path
+from typing import Literal
+
+import pymysql
+
+from .parse import extract_table_names_from_sql
+from .init import DB_PATH, logger, get_settings
+import urllib
+import polars as pl
+
+# ======================== 本地数据库 catdb ========================
+def tb_path(tb_name: str) -> Path:
+    """
+    返回指定表名 完整的本地路径
+    Parameters
+    ----------
+    tb_name: str
+       表名，路径写法: a/b/c
+    Returns
+    -------
+    pathlib.Path
+        full_abs_path: pathlib.Path
+        完整的本地绝对路径 $DB_PATH/a/b/c
+    """
+    return Path(DB_PATH, tb_name)
+
+
+def put(df, tb_name: str, partitions: list[str] | None = None):
+    """
+    将一个DataFrame写入到指定名称的表格目录中，支持分区存储。
+
+    该函数负责将给定的DataFrame（df）根据提供的表名（tb_name）写入到本地文件系统中。
+    如果指定了分区（partitions），则会按照这些分区列将数据分割存储。如果目录不存在，会自动创建目录。
+
+    Parameters
+    ----------
+    df: polars.DataFrame
+    tb_name: str
+        表的名称，用于确定存储数据的目录
+    partitions: list[str] | None
+        指定用于分区的列名列表。如果未提供，则不进行分区。
+
+    Returns
+    -------
+
+    """
+    if df is None:
+        logger.warning(f"put failed: input data is None.")
+        return
+    if df.is_empty():
+        logger.warning(f"put failed: input data is empty.")
+        return
+    tbpath = tb_path(tb_name)
+    if not tbpath.exists():
+        tbpath.mkdir(parents=True, exist_ok=True)
+    if partitions is not None:
+        df.write_parquet(tbpath, partition_by=partitions)
+    else:
+        df.write_parquet(tbpath / "data.parquet")
+
+def has(tb_name: str) -> bool:
+    """
+    判定给定的表名是否存在
+    Parameters
+    ----------
+    tb_name: str
+
+    Returns
+    -------
+
+    """
+    return tb_path(tb_name).exists()
+
+def sql(query: str, ):
+    """
+    sql 查询，从本地paquet文件中查询数据
+
+    Parameters
+    ----------
+    query: str
+        sql查询语句
+    Returns
+    -------
+
+    """
+    import polars as pl
+
+    tbs = extract_table_names_from_sql(query)
+    convertor = dict()
+    for tb in tbs:
+        db_path = tb_path(tb)
+        format_tb = f"read_parquet('{db_path}/**/*.parquet')"
+        convertor[tb] = format_tb
+    pattern = re.compile("|".join(re.escape(k) for k in convertor.keys()))
+    new_query = pattern.sub(lambda m: convertor[m.group(0)], query)
+    return pl.sql(new_query)
+
+def scan(tb: str,) -> pl.LazyFrame:
+    """polars.scan_parquet"""
+    tb = tb_path(tb)
+    return pl.scan_parquet(tb)
+
+def read_mysql(query: str, db_conf: str = "DATABASES.mysql"):
+    """
+    从MySQL数据库中读取数据。
+    Parameters
+    ----------
+    query: str
+        查询语句
+    db_conf: str
+        对应的配置 $DB_PATH/conf/settings.toml
+    Returns
+    -------
+    polars.DataFrame
+    """
+    import polars as pl
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        uri = f"mysql://{user}:{password}@{db_setting['url']}/{db_setting['db']}"
+        return pl.read_database_uri(query, uri)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute MySQL query: {e}") from e
+
+def write_mysql(df: pl.DataFrame,
+                remote_tb: str,
+                db_conf: str,
+                if_table_exists: Literal["append", "replace", "fail"]="append"):
+    """将 polars.DataFrame 写入mysql"""
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        uri = f"mysql+pymysql://{user}:{password}@{db_setting['url']}/{db_setting['db']}"
+        return df.write_database(remote_tb,
+                                 connection=uri,
+                                 if_table_exists=if_table_exists)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to write MySQL: {e}") from e
+
+def execute_mysql(sql: str, db_conf: str):
+    """执行mysql语句"""
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        url = urllib.parse.quote_plus(db_setting["url"])
+        host, port = url.split(":")
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to parse config: {e}") from e
+
+    connection = pymysql.connect(
+        host=host,
+        port=port,
+        user=user,
+        password=password,
+        database=db_setting['db']  # or extract from connection string
+    )
+    try:
+        with connection.cursor() as cursor:
+            cursor.execute(sql)
+        connection.commit()
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute MySQL: {e}") from e
+    finally:
+        connection.close()
+
+
+def read_ck(query: str, db_conf: str = "DATABASES.ck"):
+    """
+    从Clickhouse集群读取数据。
+    Parameters
+    ----------
+    query: str
+        查询语句
+    db_conf: str
+        对应的配置 $DB_PATH/conf/settings.toml
+    Returns
+    -------
+    polars.DataFrame
+    """
+    import clickhouse_df
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'urls']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+
+        with clickhouse_df.connect(db_setting['urls'], user=user, password=password):
+            return clickhouse_df.to_polars(query)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute ClickHouse query: {e}") from e

lidb-1.1.10/lidb/dataset.py

@@ -0,0 +1,310 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+# Created on 2025/10/27 14:13
+# Description:
+
+from __future__ import annotations
+
+from enum import Enum
+
+from functools import partial
+import polars.selectors as cs
+import polars as pl
+from typing import Callable, Literal
+import logair
+import xcals
+import ygo
+from .database import put, tb_path, scan, DB_PATH
+from .parse import parse_hive_partition_structure
+
+class InstrumentType(Enum):
+
+    STOCK = "Stock" # 股票
+    ETF = "ETF" #
+    CB = "ConvertibleBond" # 可转债
+
+def complete_data(fn, date, save_path, partitions):
+
+    logger = logair.get_logger(__name__)
+    try:
+        data = fn(date=date)
+        if data is None:
+            # 保存数据的逻辑在fn中实现了
+            return
+        # 剔除以 `_` 开头的列
+        data = data.select(~cs.starts_with("_"))
+        if not isinstance(data, (pl.DataFrame, pl.LazyFrame)):
+            logger.error(f"{save_path}: Result of dataset.fn must be polars.DataFrame or polars.LazyFrame.")
+            return
+        if isinstance(data, pl.LazyFrame):
+            data = data.collect()
+        cols = data.columns
+        if "date" not in cols:
+            data = data.with_columns(pl.lit(date).alias("date")).select("date", *cols)
+
+        put(data, save_path, partitions=partitions)
+    except Exception as e:
+        logger.error(f"{save_path}: Error when complete data for {date}")
+        logger.warning(e)
+
+class Dataset:
+
+    def __init__(self,
+                 fn: Callable[..., pl.DataFrame],
+                 tb: str,
+                 update_time: str = "",
+                 partitions: list[str] = None,
+                 by_asset: bool = True,
+                 by_time: bool = False):
+        """
+
+        Parameters
+        ----------
+        fn: str
+            数据集计算函数
+        tb: str
+            数据集保存表格
+        update_time: str
+            更新时间: 默认没有-实时更新，也就是可以取到当天值
+        partitions: list[str]
+            分区
+        by_asset: bool
+            是否按照标的进行分区，默认 True
+        by_time: bool
+            是否按照标的进行分区，默认 False
+        """
+        self.fn = fn
+        self.fn_params_sig = ygo.fn_signature_params(fn)
+        self._by_asset = by_asset
+        self._by_time = by_time
+        self._append_partitions = ["asset", "date"] if by_asset else ["date", ]
+        if by_time:
+            self._append_partitions.append("time")
+        if partitions is not None:
+            partitions = [k for k in partitions if k not in self._append_partitions]
+            partitions = [*partitions, *self._append_partitions]
+        else:
+            partitions = self._append_partitions
+        self.partitions = partitions
+        self._type_asset = "asset" in self.fn_params_sig
+        self.update_time = update_time
+
+        self.tb = tb
+        self.save_path = tb_path(tb)
+        fn_params = ygo.fn_params(self.fn)
+        self.fn_params = {k: v for (k, v) in fn_params}
+        self.constraints = dict()
+        for k in self.partitions[:-len(self._append_partitions)]:
+            if k in self.fn_params:
+                v = self.fn_params[k]
+                if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                    v = sorted(v)
+                self.constraints[k] = v
+                self.save_path = self.save_path / f"{k}={v}"
+
+    def is_empty(self, path) -> bool:
+        return not any(path.rglob("*.parquet"))
+
+    def __call__(self, *fn_args, **fn_kwargs):
+        # self.fn =
+        fn = partial(self.fn, *fn_args, **fn_kwargs)
+        ds = Dataset(fn=fn,
+                     tb=self.tb,
+                     partitions=self.partitions,
+                     by_asset=self._by_asset,
+                     by_time=self._by_time,
+                     update_time=self.update_time)
+        return ds
+
+    def get_value(self, date, eager: bool = True, **constraints):
+        """
+        取值: 不保证未来数据
+        Parameters
+        ----------
+        date: str
+            取值日期
+        eager: bool
+        constraints: dict
+            取值的过滤条件
+
+        Returns
+        -------
+
+        """
+        _constraints = {k: v for k, v in constraints.items() if k in self.partitions}
+        _limits = {k: v for k, v in constraints.items() if k not in self.partitions}
+        search_path = self.save_path
+        for k, v in _constraints.items():
+            if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                v = sorted(v)
+            search_path = search_path / f"{k}={v}"
+        search_path = search_path / f"date={date}"
+
+        if not self.is_empty(search_path):
+            lf = scan(search_path).cast({"date": pl.Utf8})
+            schema = lf.collect_schema()
+            _limits = {k: v for k, v in constraints.items() if schema.get(k) is not None}
+            lf = lf.filter(date=date, **_limits)
+            if not eager:
+                return lf
+            data = lf.collect()
+            if not data.is_empty():
+                return data
+        fn = self.fn
+        save_path = self.save_path
+
+        if self._type_asset:
+            if "asset" in _constraints:
+                fn = ygo.delay(self.fn)(asset=_constraints["asset"])
+        if len(self.constraints) < len(self.partitions) - len(self._append_partitions):
+            # 如果分区指定的字段没有在Dataset定义中指定，需要在get_value中指定
+            params = dict()
+            for k in self.partitions[:-len(self._append_partitions)]:
+                if k not in self.constraints:
+                    v = constraints[k]
+                    params[k] = v
+                    save_path =save_path / f"{k}={v}"
+            fn = ygo.delay(self.fn)(**params)
+        logger = logair.get_logger(__name__)
+
+        today = xcals.today()
+        now = xcals.now()
+        if (date > today) or (date == today and now < self.update_time):
+            logger.warning(f"{self.tb}: {date} is not ready, waiting for {self.update_time}")
+            return
+        complete_data(fn, date, save_path, self._append_partitions)
+
+        lf = scan(search_path).cast({"date": pl.Utf8})
+        schema = lf.collect_schema()
+        _limits = {k: v for k, v in constraints.items() if schema.get(k) is not None}
+        lf = lf.filter(date=date, **_limits)
+        if not eager:
+            return lf
+        return lf.collect()
+
+    def get_pit(self, date: str, query_time: str, eager: bool = True, **contraints):
+        """取值：如果取值时间早于更新时间，则返回上一天的值"""
+        if not self.update_time:
+            return self.get_value(date, **contraints)
+        val_date = date
+        if query_time < self.update_time:
+            val_date = xcals.shift_tradeday(date, -1)
+        return self.get_value(val_date, eager=eager, **contraints).with_columns(date=pl.lit(date),)
+
+    def get_history(self,
+                    dateList: list[str],
+                    n_jobs: int = 5,
+                    backend: Literal["threading", "multiprocessing", "loky"] = "loky",
+                    eager: bool = True,
+                    rep_asset: str = "000001", # 默认 000001
+                    **constraints):
+        """获取历史值: 不保证未来数据"""
+        _constraints = {k: v for k, v in constraints.items() if k in self.partitions}
+        search_path = self.save_path
+        for k, v in _constraints.items():
+            if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                v = sorted(v)
+            search_path = search_path / f"{k}={v}"
+        if self.is_empty(search_path):
+            # 需要补全全部数据
+            missing_dates = dateList
+        else:
+            if not self._type_asset:
+                _search_path = self.save_path
+                for k, v in _constraints.items():
+                    if k != "asset":
+                        _search_path = _search_path / f"{k}={v}"
+                    else:
+                        _search_path = _search_path / f"asset={rep_asset}"
+                hive_info = parse_hive_partition_structure(_search_path)
+            else:
+                hive_info = parse_hive_partition_structure(search_path)
+            exist_dates = hive_info["date"].to_list()
+            missing_dates = set(dateList).difference(set(exist_dates))
+            missing_dates = sorted(list(missing_dates))
+        if missing_dates:
+            fn = self.fn
+            save_path = self.save_path
+
+            if self._type_asset:
+                if "asset" in _constraints:
+                    fn = ygo.delay(self.fn)(asset=_constraints["asset"])
+
+            if len(self.constraints) < len(self.partitions) - len(self._append_partitions):
+                params = dict()
+                for k in self.partitions[:-len(self._append_partitions)]:
+                    if k not in self.constraints:
+                        v = constraints[k]
+                        params[k] = v
+                        save_path =save_path / f"{k}={v}"
+                fn = ygo.delay(self.fn)(**params)
+
+            with ygo.pool(n_jobs=n_jobs, backend=backend) as go:
+                info_path = self.save_path
+                try:
+                    info_path = info_path.relative_to(DB_PATH)
+                except:
+                    pass
+                for date in missing_dates:
+                    go.submit(complete_data, job_name=f"Completing {info_path}")(
+                        fn=fn,
+                        date=date,
+                        save_path=save_path,
+                        partitions=self._append_partitions,
+                    )
+                go.do()
+        data = scan(search_path, ).cast({"date": pl.Utf8}).filter(pl.col("date").is_in(dateList), **constraints)
+        data = data.sort("date")
+        if eager:
+            return data.collect()
+        return data
+
+def loader(ds: Dataset,
+           date_list: list[str],
+           prev_date_list: list[str],
+           prev_date_mapping: dict[str, str],
+           time: str,
+           **constraints) -> pl.LazyFrame:
+    if time < ds.update_time:
+        lf = ds.get_history(prev_date_list, eager=False, **constraints)
+    else:
+        lf = ds.get_history(date_list, eager=False, **constraints)
+    schema = lf.collect_schema()
+    include_time = schema.get("time") is not None
+    if include_time:
+        lf = lf.filter(time=time)
+    else:
+        lf = lf.with_columns(time=pl.lit(time))
+    if time < ds.update_time:
+        lf = lf.with_columns(date=pl.col("date").replace(prev_date_mapping))
+    return lf
+
+
+def load_ds(ds_list: list[Dataset],
+            beg_date: str,
+            end_date: str,
+            time: str,
+            n_jobs: int = 7,
+            backend: Literal["threading", "multiprocessing", "loky"] = "threading",
+            **constraints) -> pl.DataFrame:
+    """加载数据集：时间对齐"""
+    lfList = list()
+    date_list = xcals.get_tradingdays(beg_date, end_date)
+    beg_date, end_date = date_list[0], date_list[-1]
+    prev_date_list = xcals.get_tradingdays(xcals.shift_tradeday(beg_date, -1), xcals.shift_tradeday(end_date, -1))
+    prev_date_mapping = {prev_date: date_list[i] for i, prev_date in enumerate(prev_date_list)}
+    with ygo.pool(n_jobs=n_jobs, backend=backend) as go:
+        for ds in ds_list:
+            go.submit(loader)(ds=ds,
+                              date_list=date_list,
+                              prev_date_list=prev_date_list,
+                              prev_date_mapping=prev_date_mapping,
+                              time=time,
+                              **constraints)
+        for lf in go.do():
+            lfList.append(lf)
+    index = ("date", "time", "asset")
+    LF = pl.concat(lfList, how="align")
+    LF = LF.sort("date", "time", "asset")
+    return LF.select(*index, cs.exclude(index)).collect()
+

lidb-1.1.10/lidb/init.py

@@ -0,0 +1,42 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+# Created on 2025/7/17 14:40
+# Description:
+
+from pathlib import Path
+from dynaconf import Dynaconf
+import logair
+
+
+USERHOME = Path("~").expanduser() # 用户家目录
+NAME = "lidb"
+DB_PATH = USERHOME / NAME
+CONFIG_PATH = USERHOME / ".config" / NAME / "settings.toml"
+
+logger = logair.get_logger(NAME)
+
+
+if not CONFIG_PATH.exists():
+    try:
+        CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
+    except Exception as e:
+        logger.error(f"Failed to create settings file: {e}")
+    with open(CONFIG_PATH, "w") as f:
+        template_content = f'[global]\npath="{DB_PATH}"\n'
+    with open(CONFIG_PATH, "w") as f:
+        f.write(template_content)
+    logger.info(f"Settings file created: {CONFIG_PATH}")
+
+def get_settings():
+    try:
+        return Dynaconf(settings_files=[CONFIG_PATH])
+    except Exception as e:
+        logger.error(f"Read settings file failed: {e}")
+        return {}
+
+# 读取配置文件覆盖
+_settiings = get_settings()
+if _settiings is not None:
+    setting_db_path = _settiings.get(f"global.path", "")
+    if setting_db_path:
+        DB_PATH = Path(setting_db_path)