feat: improve search and symbols_is_namespace

- search accepts namespace as an argument
- search and symbols_is_namespace return DataFrame (default) or json

Author: chilango74

main.py

@@ -1,12 +1,3 @@
-import numpy as np
 import okama as ok
 
-asset_list = ok.AssetList(assets=['MSFT.US'], ccy='USD')
-
-print(asset_list.dividends_annual)
-
-
-x = asset_list.get_dividend_mean_growth_rate(period=20)
-# x.replace([np.inf, -np.inf], 0, inplace=True)
-
-print(f'Growth rate:', x)
+print(ok.search('aeroflot', namespace=None, response_format='frame'))

okama/api/namespaces.py

@@ -14,11 +14,17 @@ def get_namespaces():
 
 
 @lru_cache()
-def symbols_in_namespace(namespace: str = default_namespace):
+def symbols_in_namespace(namespace: str = default_namespace, response_format: str = 'frame') -> pd.DataFrame:
     string_response = API.get_symbols_in_namespace(namespace.upper())
     list_of_symbols = json.loads(string_response)
-    df = pd.DataFrame(list_of_symbols[1:], columns=list_of_symbols[0])
-    return df.astype("string", copy=False)
+    if response_format.lower() == 'frame':
+        df = pd.DataFrame(list_of_symbols[1:], columns=list_of_symbols[0])
+        return df.astype("string", copy=False)
+    elif response_format.lower() == 'json':
+        return list_of_symbols
+    else:
+        raise ValueError('response_format must be "json" or "frame"')
+
 
 
 @lru_cache()

okama/api/search.py

@@ -1,8 +1,32 @@
 import json
+from typing import Optional
+
+import pandas as pd
 
 from .api_methods import API
+from .namespaces import symbols_in_namespace
 
 
-def search(search_string: str) -> json:
+def search(search_string: str, namespace: Optional[str] = None, response_format: str = 'frame') -> json:
+    # search for string in a single namespace
+    if namespace:
+        df = symbols_in_namespace(namespace.upper())
+        condition1 = df['name'].str.contains(search_string, case=False)
+        condition2 = df['ticker'].str.contains(search_string, case=False)
+        frame_response = df[condition1 | condition2]
+        if response_format.lower() == 'frame':
+            return frame_response
+        elif response_format.lower() == 'json':
+            return frame_response.to_json(orient='records')
+        else:
+            raise ValueError('response_format must be "json" or "frame"')
+    # search for string in all namespaces
     string_response = API.search(search_string)
-    return json.loads(string_response)
+    json_response = json.loads(string_response)
+    if response_format.lower() == 'frame':
+        df = pd.DataFrame(json_response[1:], columns=json_response[0])
+        return df
+    elif response_format.lower() == 'json':
+        return json_response
+    else:
+        raise ValueError('response_format must be "json" or "frame"')

okama/common/make_asset_list.py

@@ -79,7 +79,7 @@ def __init__(
             self.assets_first_dates.update({self.inflation: Inflation(self.inflation).first_date})
             self.assets_last_dates.update({self.inflation: Inflation(self.inflation).last_date})
         self.assets_ror: pd.DataFrame = self.assets_ror[
-            self.first_date : self.last_date
+            self.first_date: self.last_date
         ]
         self.period_length: float = round(
             (self.last_date - self.first_date) / np.timedelta64(365, "D"), ndigits=1

tests/test_search.py

@@ -0,0 +1,47 @@
+"""
+Tests the search
+"""
+import json
+
+import pytest
+
+from okama.api.search import search
+
+
+def test_search_namespace_json():
+    x = search(
+        "aeroflot",
+        namespace="MOEX",
+        response_format="json",
+    )
+    assert json.loads(x)[0]['symbol'] == 'AFLT.MOEX'
+
+
+def test_search_namespace_frame():
+    x = search(
+        "aeroflot",
+        namespace="MOEX",
+        response_format="frame",
+    )
+    assert x['symbol'].values[0] == 'AFLT.MOEX'
+
+
+def test_search_all_json():
+    x = search(
+        "lkoh",
+        response_format="json",
+    )
+    assert x[1][0] == 'LKOH.MOEX'
+
+
+def test_search_all_frame():
+    x = search(
+        "lkoh",
+        response_format="frame",
+    )
+    assert x['symbol'].iloc[0] == 'LKOH.MOEX'
+
+
+def test_search_error():
+    with pytest.raises(ValueError):
+        search("arg", response_format='txt')