tickflow 0.1.0__py3-none-any.whl

tickflow/resources/klines.py

@@ -0,0 +1,460 @@
+"""K-line (OHLCV) data resources for TickFlow API."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union, overload
+
+from .._types import NOT_GIVEN, NotGiven
+from ._base import AsyncResource, SyncResource
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+    from ..generated_model import CompactKlineData, Period
+
+
+def _klines_to_dataframe(
+    data: "CompactKlineData", symbol: Optional[str] = None
+) -> "pd.DataFrame":
+    """Convert compact K-line data to a pandas DataFrame.
+
+    Parameters
+    ----------
+    data : CompactKlineData
+        Compact columnar K-line data from the API.
+    symbol : str, optional
+        Symbol code to include as a column.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with columns: timestamp, open, high, low, close, volume, amount.
+        Index is set to timestamp converted to datetime.
+    """
+    import pandas as pd
+
+    df = pd.DataFrame(
+        {
+            "timestamp": data["timestamp"],
+            "open": data["open"],
+            "high": data["high"],
+            "low": data["low"],
+            "close": data["close"],
+            "volume": data["volume"],
+            "amount": data.get("amount", [0.0] * len(data["timestamp"])),
+        }
+    )
+
+    if symbol:
+        df.insert(0, "symbol", symbol)
+
+    # Convert timestamp (milliseconds) to datetime
+    df["timestamp"] = pd.to_datetime(df["timestamp"], unit="ms")
+    df.set_index("timestamp", inplace=True)
+
+    return df
+
+
+def _batch_klines_to_dataframe(data: Dict[str, "CompactKlineData"]) -> "pd.DataFrame":
+    """Convert batch K-line data to a single pandas DataFrame.
+
+    Parameters
+    ----------
+    data : dict
+        Dictionary mapping symbol codes to compact K-line data.
+
+    Returns
+    -------
+    pd.DataFrame
+        Combined DataFrame with a 'symbol' column and MultiIndex (symbol, timestamp).
+    """
+    import pandas as pd
+
+    dfs = []
+    for symbol, kline_data in data.items():
+        df = _klines_to_dataframe(kline_data, symbol=symbol)
+        dfs.append(df)
+
+    if not dfs:
+        return pd.DataFrame()
+
+    combined = pd.concat(dfs)
+    combined.reset_index(inplace=True)
+    combined.set_index(["symbol", "timestamp"], inplace=True)
+    combined.sort_index(inplace=True)
+
+    return combined
+
+
+class Klines(SyncResource):
+    """Synchronous interface for K-line (OHLCV) data endpoints.
+
+    Supports returning data as raw dicts or pandas DataFrames.
+
+    Examples
+    --------
+    >>> client = TickFlow(api_key="your-key")
+    >>>
+    >>> # Get raw K-line data
+    >>> klines = client.klines.get("600000.SH", count=100)
+    >>>
+    >>> # Get as DataFrame
+    >>> df = client.klines.get("600000.SH", count=100, as_dataframe=True)
+    >>> print(df.head())
+    """
+
+    @overload
+    def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[False] = False,
+    ) -> "CompactKlineData": ...
+
+    @overload
+    def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[True],
+    ) -> "pd.DataFrame": ...
+
+    def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: bool = False,
+    ) -> Union["CompactKlineData", "pd.DataFrame"]:
+        """Get K-line (OHLCV) data for a single symbol.
+
+        Parameters
+        ----------
+        symbol : str
+            Symbol code (e.g., "600000.SH", "AAPL.US").
+        period : str, optional
+            K-line period. One of: "1m", "5m", "10m", "15m", "30m", "60m",
+            "4h", "1d", "1w", "1M". Defaults to "1d".
+        count : int, optional
+            Number of K-lines to return. Default 100, max 10000.
+        start_time : int, optional
+            Start timestamp in milliseconds.
+        end_time : int, optional
+            End timestamp in milliseconds.
+        as_dataframe : bool, optional
+            If True, return a pandas DataFrame. If False (default), return
+            the raw compact columnar data.
+
+        Returns
+        -------
+        CompactKlineData or pd.DataFrame
+            If as_dataframe=False: dict with keys 'timestamp', 'open', 'high',
+            'low', 'close', 'volume', 'amount' (each a list).
+            If as_dataframe=True: pandas DataFrame with datetime index.
+
+        Examples
+        --------
+        >>> # Get raw data
+        >>> data = client.klines.get("600000.SH", period="1d", count=10)
+        >>> print(data["close"][-1])  # Latest close price
+
+        >>> # Get as DataFrame for analysis
+        >>> df = client.klines.get("600000.SH", period="1d", count=100, as_dataframe=True)
+        >>> print(df.tail())
+        >>>
+        >>> # Calculate 20-day moving average
+        >>> df["ma20"] = df["close"].rolling(20).mean()
+        """
+        params = {"symbol": symbol}
+        if not isinstance(period, NotGiven):
+            params["period"] = period
+        if not isinstance(count, NotGiven):
+            params["count"] = count
+        if not isinstance(start_time, NotGiven) and start_time is not None:
+            params["start_time"] = start_time
+        if not isinstance(end_time, NotGiven) and end_time is not None:
+            params["end_time"] = end_time
+
+        response = self._client.get("/v1/klines", params=params)
+        data = response["data"]
+
+        if as_dataframe:
+            return _klines_to_dataframe(data, symbol=symbol)
+        return data
+
+    @overload
+    def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[False] = False,
+    ) -> Dict[str, "CompactKlineData"]: ...
+
+    @overload
+    def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[True],
+    ) -> "pd.DataFrame": ...
+
+    def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: bool = False,
+    ) -> Union[Dict[str, "CompactKlineData"], "pd.DataFrame"]:
+        """Get K-line data for multiple symbols in a single request.
+
+        Parameters
+        ----------
+        symbols : list of str
+            List of symbol codes.
+        period : str, optional
+            K-line period. Defaults to "1d".
+        count : int, optional
+            Number of K-lines per symbol. Default 100.
+        start_time : int, optional
+            Start timestamp in milliseconds.
+        end_time : int, optional
+            End timestamp in milliseconds.
+        as_dataframe : bool, optional
+            If True, return a combined pandas DataFrame with MultiIndex.
+            If False (default), return a dict mapping symbols to K-line data.
+
+        Returns
+        -------
+        dict or pd.DataFrame
+            If as_dataframe=False: dict mapping symbol codes to CompactKlineData.
+            If as_dataframe=True: pandas DataFrame with MultiIndex (symbol, timestamp).
+
+        Examples
+        --------
+        >>> # Get raw data for multiple symbols
+        >>> data = client.klines.batch(["600000.SH", "000001.SZ"])
+        >>> for symbol, klines in data.items():
+        ...     print(f"{symbol}: {len(klines['timestamp'])} bars")
+
+        >>> # Get as combined DataFrame
+        >>> df = client.klines.batch(["600000.SH", "000001.SZ"], as_dataframe=True)
+        >>> print(df.loc["600000.SH"].tail())
+        """
+        symbols_str = ",".join(symbols)
+        params = {"symbols": symbols_str}
+        if not isinstance(period, NotGiven):
+            params["period"] = period
+        if not isinstance(count, NotGiven):
+            params["count"] = count
+        if not isinstance(start_time, NotGiven) and start_time is not None:
+            params["start_time"] = start_time
+        if not isinstance(end_time, NotGiven) and end_time is not None:
+            params["end_time"] = end_time
+
+        response = self._client.get("/v1/klines/batch", params=params)
+        data = response["data"]
+
+        if as_dataframe:
+            return _batch_klines_to_dataframe(data)
+        return data
+
+
+class AsyncKlines(AsyncResource):
+    """Asynchronous interface for K-line (OHLCV) data endpoints.
+
+    Supports returning data as raw dicts or pandas DataFrames.
+
+    Examples
+    --------
+    >>> async with AsyncTickFlow(api_key="your-key") as client:
+    ...     # Get raw K-line data
+    ...     klines = await client.klines.get("600000.SH", count=100)
+    ...
+    ...     # Get as DataFrame
+    ...     df = await client.klines.get("600000.SH", count=100, as_dataframe=True)
+    """
+
+    @overload
+    async def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[False] = False,
+    ) -> "CompactKlineData": ...
+
+    @overload
+    async def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[True],
+    ) -> "pd.DataFrame": ...
+
+    async def get(
+        self,
+        symbol: str,
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: bool = False,
+    ) -> Union["CompactKlineData", "pd.DataFrame"]:
+        """Get K-line (OHLCV) data for a single symbol.
+
+        Parameters
+        ----------
+        symbol : str
+            Symbol code (e.g., "600000.SH", "AAPL.US").
+        period : str, optional
+            K-line period. One of: "1m", "5m", "10m", "15m", "30m", "60m",
+            "4h", "1d", "1w", "1M". Defaults to "1d".
+        count : int, optional
+            Number of K-lines to return. Default 100, max 10000.
+        start_time : int, optional
+            Start timestamp in milliseconds.
+        end_time : int, optional
+            End timestamp in milliseconds.
+        as_dataframe : bool, optional
+            If True, return a pandas DataFrame. If False (default), return
+            the raw compact columnar data.
+
+        Returns
+        -------
+        CompactKlineData or pd.DataFrame
+            If as_dataframe=False: dict with keys 'timestamp', 'open', 'high',
+            'low', 'close', 'volume', 'amount' (each a list).
+            If as_dataframe=True: pandas DataFrame with datetime index.
+
+        Examples
+        --------
+        >>> df = await client.klines.get("600000.SH", period="1d", as_dataframe=True)
+        >>> print(df.tail())
+        """
+        params = {"symbol": symbol}
+        if not isinstance(period, NotGiven):
+            params["period"] = period
+        if not isinstance(count, NotGiven):
+            params["count"] = count
+        if not isinstance(start_time, NotGiven) and start_time is not None:
+            params["start_time"] = start_time
+        if not isinstance(end_time, NotGiven) and end_time is not None:
+            params["end_time"] = end_time
+
+        response = await self._client.get("/v1/klines", params=params)
+        data = response["data"]
+
+        if as_dataframe:
+            return _klines_to_dataframe(data, symbol=symbol)
+        return data
+
+    @overload
+    async def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[False] = False,
+    ) -> Dict[str, "CompactKlineData"]: ...
+
+    @overload
+    async def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: Literal[True],
+    ) -> "pd.DataFrame": ...
+
+    async def batch(
+        self,
+        symbols: List[str],
+        *,
+        period: Union["Period", NotGiven] = NOT_GIVEN,
+        count: Union[int, NotGiven] = NOT_GIVEN,
+        start_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        end_time: Union[int, None, NotGiven] = NOT_GIVEN,
+        as_dataframe: bool = False,
+    ) -> Union[Dict[str, "CompactKlineData"], "pd.DataFrame"]:
+        """Get K-line data for multiple symbols in a single request.
+
+        Parameters
+        ----------
+        symbols : list of str
+            List of symbol codes.
+        period : str, optional
+            K-line period. Defaults to "1d".
+        count : int, optional
+            Number of K-lines per symbol. Default 100.
+        start_time : int, optional
+            Start timestamp in milliseconds.
+        end_time : int, optional
+            End timestamp in milliseconds.
+        as_dataframe : bool, optional
+            If True, return a combined pandas DataFrame with MultiIndex.
+            If False (default), return a dict mapping symbols to K-line data.
+
+        Returns
+        -------
+        dict or pd.DataFrame
+            If as_dataframe=False: dict mapping symbol codes to CompactKlineData.
+            If as_dataframe=True: pandas DataFrame with MultiIndex (symbol, timestamp).
+
+        Examples
+        --------
+        >>> df = await client.klines.batch(["600000.SH", "000001.SZ"], as_dataframe=True)
+        >>> print(df.loc["600000.SH"].tail())
+        """
+        symbols_str = ",".join(symbols)
+        params = {"symbols": symbols_str}
+        if not isinstance(period, NotGiven):
+            params["period"] = period
+        if not isinstance(count, NotGiven):
+            params["count"] = count
+        if not isinstance(start_time, NotGiven) and start_time is not None:
+            params["start_time"] = start_time
+        if not isinstance(end_time, NotGiven) and end_time is not None:
+            params["end_time"] = end_time
+
+        response = await self._client.get("/v1/klines/batch", params=params)
+        data = response["data"]
+
+        if as_dataframe:
+            return _batch_klines_to_dataframe(data)
+        return data