cjdata 0.0.2__py3-none-any.whl

cjdata/xtquant_pipeline.py

@@ -0,0 +1,315 @@
+"""Data acquisition routines backed by xtquant."""
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timedelta
+from typing import Optional, Sequence, Any, TYPE_CHECKING
+
+import pandas as pd
+from pandas.api.types import is_datetime64_any_dtype
+
+from .db import UpsertSpec, insert_ignore, upsert_rows
+from .utils import to_yyyymmdd
+
+try:  # pragma: no cover - optional dependency
+    from tqdm import tqdm  # type: ignore
+except ImportError:  # pragma: no cover
+    tqdm = None  # type: ignore
+
+xtdata: Any
+if TYPE_CHECKING:  # pragma: no cover
+    from xtquant import xtdata as _xtdata
+
+try:  # pragma: no cover - optional dependency
+    from xtquant import xtdata as _xtdata  # type: ignore
+    xtdata = _xtdata
+except ImportError:  # pragma: no cover
+    xtdata = None
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def _require_xtquant() -> None:
+    if xtdata is None:  # pragma: no cover
+        raise RuntimeError("xtquant is not installed. Install xtquant to use this feature.")
+
+
+def _to_epoch_ms(value: Any) -> int:
+    if isinstance(value, (int, float)):
+        return int(value)
+    if isinstance(value, datetime):
+        return int(value.timestamp() * 1000)
+    return int(pd.Timestamp(value).timestamp() * 1000)
+
+
+class XtQuantPipeline:
+    def __init__(self, conn, logger: Optional[logging.Logger] = None) -> None:
+        _require_xtquant()
+        self.conn = conn
+        self.logger = logger or _LOGGER
+
+    def download_trading_calendar(self, markets: Sequence[str] = ("SH", "SZ")) -> int:
+        total = 0
+        for market in markets:
+            dates = xtdata.get_trading_dates(market)
+            if not dates:
+                continue
+            rows = [(market, _to_epoch_ms(day)) for day in dates]
+            total += insert_ignore(self.conn, "trading_days", ("market", "trade_date"), rows)
+        self.conn.commit()
+        return total
+
+    def update_sector_membership(self) -> int:
+        sectors = xtdata.get_sector_list()
+        total = 0
+        for sector in sectors:
+            stocks = xtdata.get_stock_list_in_sector(sector)
+            rows = [(sector, code) for code in stocks]
+            total += insert_ignore(self.conn, "sector_stocks", ("sector_name", "stock_code"), rows)
+        self.conn.commit()
+        return total
+
+    def update_stock_basic(self, sectors: Sequence[str] = ("沪深A股", "沪深指数", "沪深基金")) -> int:
+        cursor = self.conn.execute(
+            "SELECT DISTINCT stock_code FROM sector_stocks WHERE sector_name IN ({})".format(
+                ",".join("?" for _ in sectors)
+            ),
+            tuple(sectors),
+        )
+        codes = [row[0] for row in cursor.fetchall()]
+        rows = []
+        for code in codes:
+            detail = xtdata.get_instrument_detail(code)
+            if not detail:
+                continue
+            rows.append(
+                (
+                    code,
+                    detail.get("InstrumentName"),
+                    detail.get("ExchangeID"),
+                    self._determine_board(code, detail),
+                    detail.get("OpenDate"),
+                    detail.get("TotalShares"),
+                    detail.get("CirculatingShares"),
+                )
+            )
+        spec = UpsertSpec(
+            table="stock_basic",
+            columns=(
+                "stock_code",
+                "stock_name",
+                "market",
+                "board",
+                "listed_date",
+                "total_volume",
+                "float_volume",
+            ),
+            conflict_columns=("stock_code",),
+        )
+        updated = upsert_rows(self.conn, spec, rows, update_columns=(
+            "stock_name",
+            "market",
+            "board",
+            "listed_date",
+            "total_volume",
+            "float_volume",
+        ))
+        self.conn.commit()
+        return updated
+
+    def download_daily_for_codes(
+        self,
+        codes: Sequence[str],
+        start_date: Optional[str] = None,
+        end_date: Optional[str] = None,
+        dividend_type: str = "back_ratio",
+    ) -> int:
+        total = 0
+        start = to_yyyymmdd(start_date or "20080101")
+        end = to_yyyymmdd(end_date) if end_date else ""
+        iterable = (
+            tqdm(
+                codes,
+                desc="xtquant daily",
+                total=len(codes) if hasattr(codes, "__len__") else None,
+                leave=False,
+            )
+            if tqdm
+            else codes
+        )
+        for code in iterable:
+            rows = self._download_single(code, start, end, dividend_type)
+            if not rows:
+                continue
+            spec = UpsertSpec(
+                table="daily_k_data",
+                columns=(
+                    "date",
+                    "code",
+                    "open",
+                    "high",
+                    "low",
+                    "close",
+                    "preclose",
+                    "volume",
+                    "amount",
+                    "adjustflag",
+                    "turn",
+                    "tradestatus",
+                    "pctChg",
+                    "peTTM",
+                    "pbMRQ",
+                    "psTTM",
+                    "pcfNcfTTM",
+                    "isST",
+                    "source",
+                ),
+                conflict_columns=("date", "code", "adjustflag"),
+            )
+            total += upsert_rows(
+                self.conn,
+                spec,
+                rows,
+                update_columns=(
+                    "open",
+                    "high",
+                    "low",
+                    "close",
+                    "preclose",
+                    "volume",
+                    "amount",
+                    "turn",
+                    "tradestatus",
+                    "pctChg",
+                    "peTTM",
+                    "pbMRQ",
+                    "psTTM",
+                    "pcfNcfTTM",
+                    "isST",
+                    "source",
+                ),
+            )
+        self.conn.commit()
+        return total
+
+    def default_etf_codes(self) -> list[str]:
+        cursor = self.conn.execute(
+            "SELECT DISTINCT stock_code FROM sector_stocks WHERE sector_name = ?",
+            ("沪深基金",),
+        )
+        return [row[0] for row in cursor.fetchall()]
+
+    def _download_single(
+        self,
+        stock_code: str,
+        start_date: str,
+        end_date: str,
+        dividend_type: str,
+    ):
+        start_dt = self._next_download_date(stock_code, start_date)
+        try:
+            xtdata.download_history_data(
+                stock_code=stock_code,
+                period="1d",
+                start_time=start_dt,
+                end_time=end_date,
+                incrementally=True,
+            )
+            data = xtdata.get_market_data_ex(
+                field_list=[
+                    "time",
+                    "open",
+                    "high",
+                    "low",
+                    "close",
+                    "volume",
+                    "amount",
+                    "preClose",
+                    "turn",
+                    "tradeStatus",
+                    "pctChg",
+                    "peTTM",
+                    "pbMRQ",
+                    "psTTM",
+                    "pcfNcfTTM",
+                    "isST",
+                ],
+                stock_list=[stock_code],
+                period="1d",
+                start_time=start_dt,
+                end_time=end_date,
+                dividend_type=dividend_type,
+                fill_data=True,
+            )
+        except Exception as exc:  # pragma: no cover - network error path
+            self.logger.warning("xtquant download failed for %s: %s", stock_code, exc)
+            return []
+        frame = data.get(stock_code) if data else None
+        if frame is None or frame.empty:
+            return []
+        frame = frame.copy()
+        index_series = frame.index.to_series()
+        if not is_datetime64_any_dtype(index_series):
+            index_series = pd.to_datetime(index_series, errors="coerce")
+        if index_series.isna().all() and "time" in frame.columns:
+            time_series = pd.to_datetime(frame["time"], errors="coerce")
+            if time_series.isna().all():
+                time_series = pd.to_datetime(frame["time"], unit="ms", errors="coerce")
+            index_series = time_series
+        mask = ~index_series.isna()
+        frame = frame[mask].copy()
+        index_series = index_series[mask]
+        if frame.empty:
+            return []
+        frame["date"] = index_series.dt.strftime("%Y%m%d")
+        frame["code"] = stock_code
+        frame = frame.rename(columns={"preClose": "preclose", "tradeStatus": "tradestatus"})
+        frame["adjustflag"] = 1
+        frame["source"] = "xtquant"
+        columns = [
+            "date",
+            "code",
+            "open",
+            "high",
+            "low",
+            "close",
+            "preclose",
+            "volume",
+            "amount",
+            "adjustflag",
+            "turn",
+            "tradestatus",
+            "pctChg",
+            "peTTM",
+            "pbMRQ",
+            "psTTM",
+            "pcfNcfTTM",
+            "isST",
+            "source",
+        ]
+        for column in columns:
+            if column not in frame.columns:
+                frame[column] = None
+        frame = frame[columns]
+        return [tuple(row) for row in frame.itertuples(index=False, name=None)]
+
+    def _determine_board(self, code: str, detail: dict) -> Optional[str]:
+        if code.startswith("688"):
+            return "科创板"
+        if code.startswith("300"):
+            return "创业板"
+        if code.startswith("8"):
+            return "北交所"
+        return detail.get("InstrumentStatus")
+
+    def _next_download_date(self, code: str, start_date: str) -> str:
+        cursor = self.conn.execute(
+            "SELECT MAX(date) FROM daily_k_data WHERE code=? AND adjustflag=1",
+            (code,),
+        )
+        latest = cursor.fetchone()[0]
+        if not latest:
+            return start_date
+        latest_dt = datetime.strptime(latest, "%Y%m%d") + timedelta(days=1)
+        fallback = datetime.strptime(start_date, "%Y%m%d")
+        return max(latest_dt, fallback).strftime("%Y%m%d")

cjdata-0.0.2.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: cjdata
+Version: 0.0.2
+Summary: Utilities for building and querying local stock data from xtquant and baostock sources.
+Requires-Python: >=3.10
+Requires-Dist: pandas>=2.0
+Requires-Dist: numpy>=1.23
+Provides-Extra: baostock
+Requires-Dist: baostock>=0.8.8; extra == "baostock"
+Provides-Extra: progress
+Requires-Dist: tqdm>=4.60; extra == "progress"
+Provides-Extra: talib
+Requires-Dist: TA-Lib>=0.4; extra == "talib"
+Provides-Extra: all
+Requires-Dist: baostock>=0.8.8; extra == "all"
+Requires-Dist: tqdm>=4.60; extra == "all"
+Requires-Dist: TA-Lib>=0.4; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"

cjdata-0.0.2.dist-info/RECORD

@@ -0,0 +1,15 @@
+cjdata/__init__.py,sha256=epPQD9pG7z1JfuQG8P6fegjZB6D1EbJRLqh_FuA-zMM,473
+cjdata/__main__.py,sha256=Pl4ZPieuVFtDhheztm31d7o09XEGwBLpf1nNOVBcP2E,182
+cjdata/baostock_pipeline.py,sha256=SOuBjhP6I6wJ0i01VCui7mIIvivpKrhj8lqBs-rJ7Zs,11430
+cjdata/builder.py,sha256=zcRZqXqm9vMERLVEUHz2m6sC-nptF0E-dyxFLfwn1x0,4438
+cjdata/cli.py,sha256=wYTSpIdgzzYZkAmmFW7zsdJjsCPl6gR3ij0tOJd2xBk,2757
+cjdata/db.py,sha256=iv8fGMdvru9msU3TlgprXAF567K3qdkFfn7e8M2z6QY,5019
+cjdata/local_data.py,sha256=IhopkzTYirM4iAYxZv1_d3Gl00RLEOAggZfKDv5w-ls,19723
+cjdata/utils.py,sha256=L0R1W5h98f2JZeRTkHnqRjQDgrTakLnXbOeVoYxZOKw,864
+cjdata/xtquant_pipeline.py,sha256=-W_AAA9VpNwUjd6F6lFLRC0XDXt5AXNlKP4HH6s_TPM,10662
+cjdata/__pycache__/__main__.cpython-312.pyc,sha256=6NfjvGzKVlMk2ONRLXVRQraIPba104-fdTl3swWwiNY,347
+cjdata-0.0.2.dist-info/METADATA,sha256=8_2DoZjUc2LEYib_D8OC1vlfT5pV423nxoicx4dl7lI,675
+cjdata-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cjdata-0.0.2.dist-info/entry_points.txt,sha256=87lnCH8cJY_xs2jpPfXF5JA5memlJSdwVc9RFdzTyvs,43
+cjdata-0.0.2.dist-info/top_level.txt,sha256=SMBWvkmEOUn4wQORzUpIJheNRmSgSDCrwzd7112afzc,7
+cjdata-0.0.2.dist-info/RECORD,,

cjdata-0.0.2.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

cjdata-0.0.2.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+cjdata = cjdata.cli:main

cjdata-0.0.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+cjdata