Support equity_search

Author: shugaoye

openbb_akshare/__init__.py

@@ -13,6 +13,7 @@
 from openbb_akshare.models.equity_historical import AKShareEquityHistoricalFetcher
 from openbb_akshare.models.equity_profile import AKShareEquityProfileFetcher
 from openbb_akshare.models.equity_screener import AKShareEquityScreenerFetcher
+from openbb_akshare.models.equity_search import AKShareEquitySearchFetcher
 from openbb_akshare.models.historical_dividends import AKShareHistoricalDividendsFetcher
 from openbb_akshare.models.income_statement import AKShareIncomeStatementFetcher
 from openbb_akshare.models.key_metrics import AKShareKeyMetricsFetcher
@@ -38,6 +39,7 @@
         "EquityHistorical": AKShareEquityHistoricalFetcher,
         "EquityInfo": AKShareEquityProfileFetcher,
         "EquityScreener": AKShareEquityScreenerFetcher,
+        "EquitySearch": AKShareEquitySearchFetcher,
         "HistoricalDividends": AKShareHistoricalDividendsFetcher,
         "IncomeStatement": AKShareIncomeStatementFetcher,
         "KeyMetrics": AKShareKeyMetricsFetcher,

openbb_akshare/models/equity_search.py

@@ -0,0 +1,84 @@
+"""AKShare Equity Search Model."""
+
+from typing import Any, Dict, List, Optional
+
+from openbb_core.provider.abstract.fetcher import Fetcher
+from openbb_core.provider.standard_models.equity_search import (
+    EquitySearchData,
+    EquitySearchQueryParams,
+)
+from openbb_core.provider.utils.descriptions import (
+    DATA_DESCRIPTIONS,
+    QUERY_DESCRIPTIONS,
+)
+from pydantic import Field
+
+import logging
+from openbb_akshare import project_name
+from mysharelib.tools import setup_logger
+
+setup_logger(project_name)
+logger = logging.getLogger(__name__)
+
+
+class AKShareEquitySearchQueryParams(EquitySearchQueryParams):
+    """AKShare Equity Search Query.
+    """
+
+    use_cache: bool = Field(
+        default=True,
+        description="Whether to use a cached request. The quote is cached for one hour.",
+    )
+    limit: Optional[int] = Field(
+        default=10000,
+        description=QUERY_DESCRIPTIONS.get("limit", ""),
+    )
+
+
+class AKShareEquitySearchData(EquitySearchData):
+    """AKShare Equity Search Data."""
+
+
+class AKShareEquitySearchFetcher(
+    Fetcher[
+        AKShareEquitySearchQueryParams,
+        List[AKShareEquitySearchData],
+    ]
+):
+    """Transform the query, extract and transform the data from the AKShare endpoints."""
+
+    @staticmethod
+    def transform_query(params: Dict[str, Any]) -> AKShareEquitySearchQueryParams:
+        """Transform the query."""
+        return AKShareEquitySearchQueryParams(**params)
+
+    @staticmethod
+    async def aextract_data(
+        query: AKShareEquitySearchQueryParams,  # pylint: disable=unused-argument
+        credentials: Optional[Dict[str, str]],
+        **kwargs: Any,
+    ) -> List[Dict]:
+        """Return the raw data from the AKShare endpoint."""
+
+        from openbb_akshare.utils.ak_equity_search import get_symbols
+        api_key = credentials.get("akshare_api_key") if credentials else ""
+        data = get_symbols(query.use_cache, api_key=api_key)
+        if query.limit: data = data.head(query.limit)
+
+        return data.to_dict(orient="records")
+
+    @staticmethod
+    def transform_data(
+        query: AKShareEquitySearchQueryParams, data: Dict, **kwargs: Any
+    ) -> List[AKShareEquitySearchData]:
+        """Transform the data to the standard format."""
+
+        if query.query:
+            filtered = [
+                d for d in data
+                if query.query in d.get('name', '') or query.query in d.get('symbol', '')
+            ]
+            logger.info(f"Searching for {query.query} and found {len(filtered)} results.")
+            return [AKShareEquitySearchData.model_validate(d) for d in filtered]
+
+        return [AKShareEquitySearchData.model_validate(d) for d in data]

openbb_akshare/utils/ak_equity_search.py

@@ -0,0 +1,44 @@
+import logging
+import pandas as pd
+from mysharelib.tools import setup_logger, get_exchange
+from mysharelib.table_cache import TableCache
+from openbb_akshare import project_name
+
+setup_logger(project_name)
+logger = logging.getLogger(__name__)
+
+TABLE_SCHEMA = {
+    "symbol": "TEXT",               # Symbol (e.g. 600000)
+    "name": "TEXT",                 # Short name of the asset
+    "exchange": "TEXT"              # Exchange code (e.g. SSE/SZSE)
+}
+
+def get_symbols_df() -> pd.DataFrame:
+    import akshare as ak
+
+    stock_info_a_code_name_df = ak.stock_info_a_code_name()
+    stock_info_a_code_name_df.rename(columns={"code": "symbol"}, inplace=True)
+    stock_info_a_code_name_df['exchange'] = [
+        get_exchange(code)
+        for code in stock_info_a_code_name_df['symbol']
+        ]
+    # Get symbols for HK market
+    hk_stock_list_df = ak.stock_hk_spot_em()
+    stock_info_hk_code_name_df = hk_stock_list_df[["代码","名称"]].copy()
+    stock_info_hk_code_name_df.rename(columns={"代码": "symbol", "名称": "name"}, inplace=True)
+    stock_info_hk_code_name_df["exchange"] = "HKEX"
+    
+    return pd.concat([stock_info_a_code_name_df, stock_info_hk_code_name_df], ignore_index=True) 
+
+def get_symbols(use_cache: bool = True, api_key : str = "") -> pd.DataFrame:
+    cache = TableCache(TABLE_SCHEMA, project=project_name, table_name="symbols", primary_key="symbol")
+    if use_cache:
+        data = cache.read_dataframe()
+        if not data.empty:
+            logger.info(f"Loading symbols from {project_name} cache...")
+            return data
+
+    logger.info(f"Generating symbols for {project_name} ...")
+    data = get_symbols_df()
+    cache.write_dataframe(data)
+    return data

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "openbb-akshare"
-version = "0.9.23"
+version = "1.0.5"
 description = "AKShare extension for OpenBB"
 authors = ["Roger Ye <shugaoye@yahoo.com>"]
 license = "AGPL-3.0-only"
@@ -11,7 +11,7 @@ packages = [{ include = "openbb_akshare" }]
 python = ">=3.9.21,<3.13"
 akshare = "1.17.39"
 openbb-core = { version = "^1.4.7" }
-mysharelib = "^0.9.0"
+mysharelib = "^1.0.1"
 ipywidgets = "^8.1.7"
 
 [tool.poetry.group.dev.dependencies]