tickflow 0.1.0.dev0__py3-none-any.whl → 0.1.0.dev1__py3-none-any.whl

tickflow/resources/instruments.py

@@ -0,0 +1,180 @@
+"""Instrument resources for TickFlow API."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional, Union, overload
+
+from ._base import AsyncResource, SyncResource
+
+if TYPE_CHECKING:
+    from ..generated_model import Instrument, InstrumentType
+
+
+class Instruments(SyncResource):
+    """Synchronous interface for instrument endpoints.
+
+    Examples
+    --------
+    >>> client = TickFlow(api_key="your-key")
+    >>> inst = client.instruments.get("600000.SH")
+    >>> print(f"{inst['symbol']}: {inst['name']}")
+    """
+
+    @overload
+    def get(self, symbol: str) -> "Instrument": ...
+
+    @overload
+    def get(self, symbol: List[str]) -> List["Instrument"]: ...
+
+    def get(
+        self, symbol: Union[str, List[str]]
+    ) -> Union["Instrument", List["Instrument"]]:
+        """Get metadata for one or more instruments.
+
+        Parameters
+        ----------
+        symbol : str or list of str
+            Instrument code(s). Can be a single symbol string or a list of symbols.
+
+        Returns
+        -------
+        Instrument or list of Instrument
+            If a single symbol is provided, returns a single Instrument dict.
+            If a list is provided, returns a list of Instrument dicts.
+
+            Each Instrument contains:
+            - symbol: Full symbol code (e.g., "600000.SH")
+            - code: Exchange-specific code (e.g., "600000")
+            - exchange: Exchange code (e.g., "SH")
+            - region: Region code (e.g., "CN")
+            - name: Instrument name
+            - instrument_type: Type (stock, etf, index, etc.)
+            - ext: Market-specific extension data
+
+        Examples
+        --------
+        >>> # Single instrument
+        >>> inst = client.instruments.get("600000.SH")
+        >>> print(inst['name'])
+
+        >>> # Multiple instruments
+        >>> insts = client.instruments.get(["600000.SH", "AAPL.US"])
+        >>> for i in insts:
+        ...     print(f"{i['symbol']}: {i['name']}")
+        """
+        if isinstance(symbol, str):
+            response = self._client.get("/v1/instruments", params={"symbols": symbol})
+            data = response["data"]
+            return data[0] if data else {}
+        else:
+            # Use POST for batch queries
+            response = self._client.post("/v1/instruments", json={"symbols": symbol})
+            return response["data"]
+
+    def batch(self, symbols: List[str]) -> List["Instrument"]:
+        """Get metadata for multiple instruments.
+
+        This method uses POST to handle large batches without URL length limits.
+
+        Parameters
+        ----------
+        symbols : list of str
+            List of symbol codes (up to 1000).
+
+        Returns
+        -------
+        list of Instrument
+            List of instrument metadata dicts.
+
+        Examples
+        --------
+        >>> insts = client.instruments.batch(["600000.SH", "000001.SZ", "AAPL.US"])
+        >>> for i in insts:
+        ...     print(f"{i['symbol']}: {i['name']}")
+        """
+        response = self._client.post("/v1/instruments", json={"symbols": symbols})
+        return response["data"]
+
+
+class AsyncInstruments(AsyncResource):
+    """Asynchronous interface for instrument endpoints.
+
+    Examples
+    --------
+    >>> async with AsyncTickFlow(api_key="your-key") as client:
+    ...     inst = await client.instruments.get("600000.SH")
+    """
+
+    @overload
+    async def get(self, symbol: str) -> "Instrument": ...
+
+    @overload
+    async def get(self, symbol: List[str]) -> List["Instrument"]: ...
+
+    async def get(
+        self, symbol: Union[str, List[str]]
+    ) -> Union["Instrument", List["Instrument"]]:
+        """Get metadata for one or more instruments.
+
+        Parameters
+        ----------
+        symbol : str or list of str
+            Instrument code(s). Can be a single symbol string or a list of symbols.
+
+        Returns
+        -------
+        Instrument or list of Instrument
+            If a single symbol is provided, returns a single Instrument dict.
+            If a list is provided, returns a list of Instrument dicts.
+
+            Each Instrument contains:
+            - symbol: Full symbol code (e.g., "600000.SH")
+            - code: Exchange-specific code (e.g., "600000")
+            - exchange: Exchange code (e.g., "SH")
+            - region: Region code (e.g., "CN")
+            - name: Instrument name
+            - instrument_type: Type (stock, etf, index, etc.)
+            - ext: Market-specific extension data
+
+        Examples
+        --------
+        >>> # Single instrument
+        >>> inst = await client.instruments.get("600000.SH")
+        >>> print(inst['name'])
+
+        >>> # Multiple instruments
+        >>> insts = await client.instruments.get(["600000.SH", "AAPL.US"])
+        """
+        if isinstance(symbol, str):
+            response = await self._client.get(
+                "/v1/instruments", params={"symbols": symbol}
+            )
+            data = response["data"]
+            return data[0] if data else {}
+        else:
+            response = await self._client.post(
+                "/v1/instruments", json={"symbols": symbol}
+            )
+            return response["data"]
+
+    async def batch(self, symbols: List[str]) -> List["Instrument"]:
+        """Get metadata for multiple instruments.
+
+        This method uses POST to handle large batches without URL length limits.
+
+        Parameters
+        ----------
+        symbols : list of str
+            List of symbol codes (up to 1000).
+
+        Returns
+        -------
+        list of Instrument
+            List of instrument metadata dicts.
+
+        Examples
+        --------
+        >>> insts = await client.instruments.batch(["600000.SH", "000001.SZ", "AAPL.US"])
+        """
+        response = await self._client.post("/v1/instruments", json={"symbols": symbols})
+        return response["data"]

tickflow/resources/klines.py

@@ -2,7 +2,19 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union, overload
+import asyncio
+import concurrent.futures
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    overload,
+)
 
 from .._types import NOT_GIVEN, NotGiven
 from ._base import AsyncResource, SyncResource
@@ -12,6 +24,9 @@ if TYPE_CHECKING:
 
     from ..generated_model import CompactKlineData, Period
 
+# Maximum symbols per batch request (API limit)
+MAX_SYMBOLS_PER_BATCH = 100
+
 
 def _klines_to_dataframe(
     data: "CompactKlineData", symbol: Optional[str] = None
@@ -86,6 +101,52 @@ def _batch_klines_to_dataframe(data: Dict[str, "CompactKlineData"]) -> "pd.DataF
     return combined
 
 
+def _chunk_list(lst: List[str], chunk_size: int) -> List[List[str]]:
+    """Split a list into chunks of specified size.
+
+    Parameters
+    ----------
+    lst : list
+        The list to chunk.
+    chunk_size : int
+        Maximum size of each chunk.
+
+    Returns
+    -------
+    list of list
+        List of chunks.
+    """
+    return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
+def _get_progress_bar(total: int, desc: str, show_progress: bool):
+    """Get a progress bar iterator or a simple range.
+
+    Parameters
+    ----------
+    total : int
+        Total number of items.
+    desc : str
+        Description for the progress bar.
+    show_progress : bool
+        Whether to show the progress bar.
+
+    Returns
+    -------
+    iterator
+        tqdm progress bar or range.
+    """
+    if show_progress:
+        try:
+            from tqdm.auto import tqdm
+
+            return tqdm(total=total, desc=desc, leave=False)
+        except ImportError:
+            # tqdm not installed, fall back to no progress bar
+            pass
+    return None
+
+
 class Klines(SyncResource):
     """Synchronous interface for K-line (OHLCV) data endpoints.
 
@@ -176,7 +237,7 @@ class Klines(SyncResource):
         >>> # Calculate 20-day moving average
         >>> df["ma20"] = df["close"].rolling(20).mean()
         """
-        params = {"symbol": symbol}
+        params: Dict[str, Any] = {"symbol": symbol}
         if not isinstance(period, NotGiven):
             params["period"] = period
         if not isinstance(count, NotGiven):
@@ -193,6 +254,28 @@ class Klines(SyncResource):
             return _klines_to_dataframe(data, symbol=symbol)
         return data
 
+    def _fetch_batch_chunk(
+        self,
+        symbols: List[str],
+        params: Dict[str, Any],
+    ) -> Tuple[Dict[str, "CompactKlineData"], List[Tuple[str, Exception]]]:
+        """Fetch a single batch chunk.
+
+        Returns
+        -------
+        tuple
+            (data dict, list of (symbol, error) for failed symbols)
+        """
+        symbols_str = ",".join(symbols)
+        chunk_params = {**params, "symbols": symbols_str}
+
+        try:
+            response = self._client.get("/v1/klines/batch", params=chunk_params)
+            return response["data"], []
+        except Exception as e:
+            # Return empty data and record the error for all symbols in this chunk
+            return {}, [(s, e) for s in symbols]
+
     @overload
     def batch(
         self,
@@ -203,6 +286,8 @@ class Klines(SyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: Literal[False] = False,
+        show_progress: bool = False,
+        max_workers: int = 5,
     ) -> Dict[str, "CompactKlineData"]: ...
 
     @overload
@@ -215,6 +300,8 @@ class Klines(SyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: Literal[True],
+        show_progress: bool = False,
+        max_workers: int = 5,
     ) -> "pd.DataFrame": ...
 
     def batch(
@@ -226,13 +313,19 @@ class Klines(SyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: bool = False,
+        show_progress: bool = False,
+        max_workers: int = 5,
     ) -> Union[Dict[str, "CompactKlineData"], "pd.DataFrame"]:
-        """Get K-line data for multiple symbols in a single request.
+        """Get K-line data for multiple symbols in batched concurrent requests.
+
+        This method automatically handles the API limit of 100 symbols per request
+        by splitting into chunks and fetching concurrently. Failed chunks don't
+        affect other chunks - partial results are returned.
 
         Parameters
         ----------
         symbols : list of str
-            List of symbol codes.
+            List of symbol codes. Can exceed 100 - will be automatically chunked.
         period : str, optional
             K-line period. Defaults to "1d".
         count : int, optional
@@ -244,6 +337,10 @@ class Klines(SyncResource):
         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.
+        show_progress : bool, optional
+            If True, display a progress bar (requires tqdm). Default False.
+        max_workers : int, optional
+            Maximum number of concurrent requests. Default 5.
 
         Returns
         -------
@@ -251,19 +348,25 @@ class Klines(SyncResource):
             If as_dataframe=False: dict mapping symbol codes to CompactKlineData.
             If as_dataframe=True: pandas DataFrame with MultiIndex (symbol, timestamp).
 
+        Notes
+        -----
+        - The API limits batch requests to 100 symbols. This method automatically
+          splits larger requests into chunks.
+        - Failed chunks are logged but don't cause the entire request to fail.
+        - Use `show_progress=True` for large requests to see progress.
+
         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())
+        >>> # Get data for 500 symbols with progress bar
+        >>> symbols = client.exchanges.get_symbols("SH")[:500]
+        >>> df = client.klines.batch(symbols, as_dataframe=True, show_progress=True)
+        >>> print(f"Got data for {len(df.index.get_level_values('symbol').unique())} symbols")
         """
-        symbols_str = ",".join(symbols)
-        params = {"symbols": symbols_str}
+        if not symbols:
+            return {} if not as_dataframe else _batch_klines_to_dataframe({})
+
+        # Build base params
+        params: Dict[str, Any] = {}
         if not isinstance(period, NotGiven):
             params["period"] = period
         if not isinstance(count, NotGiven):
@@ -273,12 +376,56 @@ class Klines(SyncResource):
         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"]
+        # Split symbols into chunks
+        chunks = _chunk_list(symbols, MAX_SYMBOLS_PER_BATCH)
+
+        # If only one chunk, no need for concurrency
+        if len(chunks) == 1:
+            data, errors = self._fetch_batch_chunk(chunks[0], params)
+            if as_dataframe:
+                return _batch_klines_to_dataframe(data)
+            return data
+
+        # Setup progress bar
+        pbar = _get_progress_bar(len(chunks), "Fetching K-lines", show_progress)
+
+        # Fetch chunks concurrently
+        all_data: Dict[str, "CompactKlineData"] = {}
+        all_errors: List[Tuple[str, Exception]] = []
+
+        try:
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=max_workers
+            ) as executor:
+                futures = {
+                    executor.submit(self._fetch_batch_chunk, chunk, params): chunk
+                    for chunk in chunks
+                }
+
+                for future in concurrent.futures.as_completed(futures):
+                    try:
+                        data, errors = future.result()
+                        all_data.update(data)
+                        all_errors.extend(errors)
+                    except Exception as e:
+                        chunk = futures[future]
+                        all_errors.extend((s, e) for s in chunk)
+
+                    if pbar:
+                        pbar.update(1)
+        finally:
+            if pbar:
+                pbar.close()
+
+        # Log errors if any (could use logging module in production)
+        if all_errors:
+            error_symbols = [s for s, _ in all_errors]
+            # Silently continue - partial results are better than no results
+            # Users can check if their symbols are in the result
 
         if as_dataframe:
-            return _batch_klines_to_dataframe(data)
-        return data
+            return _batch_klines_to_dataframe(all_data)
+        return all_data
 
 
 class AsyncKlines(AsyncResource):
@@ -361,7 +508,7 @@ class AsyncKlines(AsyncResource):
         >>> df = await client.klines.get("600000.SH", period="1d", as_dataframe=True)
         >>> print(df.tail())
         """
-        params = {"symbol": symbol}
+        params: Dict[str, Any] = {"symbol": symbol}
         if not isinstance(period, NotGiven):
             params["period"] = period
         if not isinstance(count, NotGiven):
@@ -378,6 +525,27 @@ class AsyncKlines(AsyncResource):
             return _klines_to_dataframe(data, symbol=symbol)
         return data
 
+    async def _fetch_batch_chunk(
+        self,
+        symbols: List[str],
+        params: Dict[str, Any],
+    ) -> Tuple[Dict[str, "CompactKlineData"], List[Tuple[str, Exception]]]:
+        """Fetch a single batch chunk asynchronously.
+
+        Returns
+        -------
+        tuple
+            (data dict, list of (symbol, error) for failed symbols)
+        """
+        symbols_str = ",".join(symbols)
+        chunk_params = {**params, "symbols": symbols_str}
+
+        try:
+            response = await self._client.get("/v1/klines/batch", params=chunk_params)
+            return response["data"], []
+        except Exception as e:
+            return {}, [(s, e) for s in symbols]
+
     @overload
     async def batch(
         self,
@@ -388,6 +556,8 @@ class AsyncKlines(AsyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: Literal[False] = False,
+        show_progress: bool = False,
+        max_concurrency: int = 5,
     ) -> Dict[str, "CompactKlineData"]: ...
 
     @overload
@@ -400,6 +570,8 @@ class AsyncKlines(AsyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: Literal[True],
+        show_progress: bool = False,
+        max_concurrency: int = 5,
     ) -> "pd.DataFrame": ...
 
     async def batch(
@@ -411,13 +583,19 @@ class AsyncKlines(AsyncResource):
         start_time: Union[int, None, NotGiven] = NOT_GIVEN,
         end_time: Union[int, None, NotGiven] = NOT_GIVEN,
         as_dataframe: bool = False,
+        show_progress: bool = False,
+        max_concurrency: int = 5,
     ) -> Union[Dict[str, "CompactKlineData"], "pd.DataFrame"]:
-        """Get K-line data for multiple symbols in a single request.
+        """Get K-line data for multiple symbols in batched concurrent requests.
+
+        This method automatically handles the API limit of 100 symbols per request
+        by splitting into chunks and fetching concurrently. Failed chunks don't
+        affect other chunks - partial results are returned.
 
         Parameters
         ----------
         symbols : list of str
-            List of symbol codes.
+            List of symbol codes. Can exceed 100 - will be automatically chunked.
         period : str, optional
             K-line period. Defaults to "1d".
         count : int, optional
@@ -429,6 +607,10 @@ class AsyncKlines(AsyncResource):
         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.
+        show_progress : bool, optional
+            If True, display a progress bar (requires tqdm). Default False.
+        max_concurrency : int, optional
+            Maximum number of concurrent requests. Default 5.
 
         Returns
         -------
@@ -436,13 +618,24 @@ class AsyncKlines(AsyncResource):
             If as_dataframe=False: dict mapping symbol codes to CompactKlineData.
             If as_dataframe=True: pandas DataFrame with MultiIndex (symbol, timestamp).
 
+        Notes
+        -----
+        - The API limits batch requests to 100 symbols. This method automatically
+          splits larger requests into chunks.
+        - Failed chunks are logged but don't cause the entire request to fail.
+        - Use `show_progress=True` for large requests to see progress.
+
         Examples
         --------
-        >>> df = await client.klines.batch(["600000.SH", "000001.SZ"], as_dataframe=True)
-        >>> print(df.loc["600000.SH"].tail())
+        >>> # Get data for many symbols with progress bar
+        >>> symbols = await client.exchanges.get_symbols("SH")
+        >>> df = await client.klines.batch(symbols[:500], as_dataframe=True, show_progress=True)
         """
-        symbols_str = ",".join(symbols)
-        params = {"symbols": symbols_str}
+        if not symbols:
+            return {} if not as_dataframe else _batch_klines_to_dataframe({})
+
+        # Build base params
+        params: Dict[str, Any] = {}
         if not isinstance(period, NotGiven):
             params["period"] = period
         if not isinstance(count, NotGiven):
@@ -452,9 +645,50 @@ class AsyncKlines(AsyncResource):
         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"]
+        # Split symbols into chunks
+        chunks = _chunk_list(symbols, MAX_SYMBOLS_PER_BATCH)
+
+        # If only one chunk, no need for semaphore
+        if len(chunks) == 1:
+            data, errors = await self._fetch_batch_chunk(chunks[0], params)
+            if as_dataframe:
+                return _batch_klines_to_dataframe(data)
+            return data
+
+        # Setup progress bar
+        pbar = _get_progress_bar(len(chunks), "Fetching K-lines", show_progress)
+
+        # Use semaphore to limit concurrency
+        semaphore = asyncio.Semaphore(max_concurrency)
+
+        async def fetch_with_semaphore(
+            chunk: List[str],
+        ) -> Tuple[Dict[str, "CompactKlineData"], List[Tuple[str, Exception]]]:
+            async with semaphore:
+                result = await self._fetch_batch_chunk(chunk, params)
+                if pbar:
+                    pbar.update(1)
+                return result
+
+        # Fetch all chunks concurrently
+        all_data: Dict[str, "CompactKlineData"] = {}
+        all_errors: List[Tuple[str, Exception]] = []
+
+        try:
+            tasks = [fetch_with_semaphore(chunk) for chunk in chunks]
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    all_errors.extend((s, result) for s in chunks[i])
+                else:
+                    data, errors = result
+                    all_data.update(data)
+                    all_errors.extend(errors)
+        finally:
+            if pbar:
+                pbar.close()
 
         if as_dataframe:
-            return _batch_klines_to_dataframe(data)
-        return data
+            return _batch_klines_to_dataframe(all_data)
+        return all_data