import warnings
from typing import Optional, overload
from typing_extensions import Union, deprecated
from bigdata_client.api.knowledge_graph import (
AutosuggestedSavedSearch,
AutosuggestRequests,
AutosuggestResponse,
ByIdsRequest,
EntityTypes,
KnowledgeGraphTypes,
)
from bigdata_client.connection_protocol import BigdataConnectionProtocol
from bigdata_client.constants import DEPRECATED_WARNING_AUTOSUGGEST
from bigdata_client.enum_utils import StrEnum
from bigdata_client.models.entities import (
Company,
Concept,
Etf,
MacroEntity,
Organization,
Person,
Place,
Product,
)
from bigdata_client.models.languages import Language
from bigdata_client.models.sources import Source
from bigdata_client.models.topics import Topic
from bigdata_client.query_type import QueryType
class FilterAnalyticCategory(StrEnum):
"""Categories used for filtering Knowledge Graph results"""
COMPANIES = "Companies"
CONCEPTS = "Concepts"
ORGANIZATIONS = "Organizations"
PEOPLE = "People"
PLACES = "Places"
PRODUCTS = "Products"
TOPICS = "Topics"
SOURCES = "Sources"
ETF = "ETFs"
def to_expected_model(self):
# Old async autosuggest doesn't support filtering
# so it is done by removing every result except the expected model
return {
FilterAnalyticCategory.COMPANIES: Company,
FilterAnalyticCategory.CONCEPTS: Concept,
FilterAnalyticCategory.ORGANIZATIONS: Organization,
FilterAnalyticCategory.PEOPLE: Person,
FilterAnalyticCategory.PLACES: Place,
FilterAnalyticCategory.PRODUCTS: Product,
FilterAnalyticCategory.TOPICS: Topic,
FilterAnalyticCategory.SOURCES: Source,
FilterAnalyticCategory.ETF: Etf,
}.get(self)
[docs]
class KnowledgeGraph:
"""For finding entities, sources and topics"""
[docs]
def __init__(self, api_connection: BigdataConnectionProtocol):
self._api = api_connection
def _autosuggest(
self,
value: str,
limit: int,
categories: Optional[list[FilterAnalyticCategory]] = None,
) -> list[KnowledgeGraphTypes]:
return self._api.autosuggest(value, limit=limit, categories=categories)
def _autosuggest_async(
self,
values: list[str],
limit: int,
) -> AutosuggestResponse:
return self._api.autosuggest_async(
AutosuggestRequests(root=values), limit=limit
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def autosuggest(
self, values: list[str], /, limit: int = 20
) -> dict[str, list[KnowledgeGraphTypes]]: ...
@overload
def autosuggest(
self, value: str, /, limit: int = 20
) -> list[KnowledgeGraphTypes]: ...
[docs]
def autosuggest(
self, values: Union[list[str], str], limit=20
) -> Union[dict[str, list[KnowledgeGraphTypes]], list[KnowledgeGraphTypes]]:
"""
Searches for entities, sources, topics, searches and watchlists
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result
Returns:
Dict with the searched terms as keys each with a list of results.
"""
if isinstance(values, list):
# Decorator not displaying the msg
warnings.warn(
DEPRECATED_WARNING_AUTOSUGGEST, DeprecationWarning, stacklevel=2
)
api_response = self._autosuggest_async(values, limit)
# Exclude macros and saved searches from response
only_supported_entities = self._exclude_models(
api_response, models=(MacroEntity, AutosuggestedSavedSearch)
)
return dict(only_supported_entities.root.items())
if isinstance(values, str):
api_response = self._autosuggest(value=values, limit=limit)
# Exclude macros and saved searches from response
only_supported_entities = self._exclude_models(
api_response, models=(MacroEntity, AutosuggestedSavedSearch)
)
return only_supported_entities
raise TypeError(f"Expected list or str for @values, found {type(values)}")
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_concepts(
self, values: list[str], /, limit=20
) -> dict[str, list[Concept]]: ...
@overload
def find_concepts(self, values: str, /, limit=20) -> list[Concept]: ...
[docs]
def find_concepts(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Concept]], list[Concept]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a concept.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.CONCEPTS, values=values, limit=limit
)
[docs]
def get_companies_by_isin(self, isins: list[str]) -> list[Optional[Company]]:
"""Retrieve a list of companies by their ISIN
Args:
isins (list[str]): ISIN list
Returns:
list[Optional[Company]]: List of companies in the same order as original @isin list, or None if was not found
"""
return self._api.get_companies_by_isin(isins)
[docs]
def get_companies_by_cusip(self, cusips: list[str]) -> list[Optional[Company]]:
"""Retrieve a list of companies by their CUSIP
Args:
cusips (list[str]): CUSIP list
Returns:
list[Optional[Company]]: List of companies in the same order as original @cusip list, or None if was not found
"""
return self._api.get_companies_by_cusip(cusips)
[docs]
def get_companies_by_sedol(self, sedols: list[str]) -> list[Optional[Company]]:
"""Retrieve a list of companies by their SEDOL
Args:
sedols (list[str]): SEDOL list
Returns:
list[Optional[Company]]: List of companies in the same order as original @sedol list, or None if was not found
"""
return self._api.get_companies_by_sedol(sedols)
[docs]
def get_companies_by_listing(self, listings: list[str]) -> list[Optional[Company]]:
"""Retrieve a list of companies by their listing
Args:
listings (list[str]): listing list
Returns:
list[Optional[Company]]: List of companies in the same order as original @listing list, or None if was not found
"""
return self._api.get_companies_by_listing(listings)
[docs]
def find_etfs(self, value: str, /, limit=20):
"""
Searches for value in the Knowledge Graph and filters out anything that is not a ETF.
Args:
value: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.ETF,
values=value,
limit=limit,
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_companies(
self, values: list[str], /, limit=20
) -> dict[str, list[Company]]: ...
@overload
def find_companies(self, values: str, /, limit=20) -> list[Company]: ...
[docs]
def find_companies(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Company]], list[Company]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a company.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.COMPANIES,
values=values,
limit=limit,
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_people(self, values: list[str], /, limit=20): ...
@overload
def find_people(self, values: str, /, limit=20) -> list[Person]: ...
[docs]
def find_people(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Person]], list[Person]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a person.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.PEOPLE, values=values, limit=limit
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_places(self, values: list[str], /, limit=20) -> dict[str, list[Place]]: ...
@overload
def find_places(self, values: str, /, limit=20) -> list[Place]: ...
[docs]
def find_places(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Place]], list[Place]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a place.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.PLACES, values=values, limit=limit
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_organizations(
self, values: list[str], /, limit=20
) -> dict[str, list[Organization]]: ...
@overload
def find_organizations(self, values: str, /, limit=20) -> list[Organization]: ...
[docs]
def find_organizations(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Organization]], list[Organization]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not an organization.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.ORGANIZATIONS,
values=values,
limit=limit,
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_products(
self, values: list[str], /, limit=20
) -> dict[str, list[Product]]: ...
@overload
def find_products(self, values: str, /, limit=20) -> list[Product]: ...
[docs]
def find_products(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Product]], list[Product]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a product.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.PRODUCTS, values=values, limit=limit
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_sources(
self, values: list[str], /, limit=20
) -> dict[str, list[Source]]: ...
@overload
def find_sources(self, values: str, /, limit=20) -> list[Source]: ...
[docs]
def find_sources(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Source]], list[Source]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a source.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.SOURCES, values=values, limit=limit
)
@overload
@deprecated(DEPRECATED_WARNING_AUTOSUGGEST)
def find_topics(self, values: list[str], /, limit=20) -> dict[str, list[Topic]]: ...
@overload
def find_topics(self, value: str, /, limit=20) -> list[Topic]: ...
[docs]
def find_topics(
self, values: Union[list[str], str], /, limit=20
) -> Union[dict[str, list[Topic]], list[Topic]]:
"""
Searches for values in the Knowledge Graph and filters out anything that is not a topic.
-------------------
Overloaded method
-------------------
* Implementation 1
Args:
values: Searched item (str)
limit: Upper limit for each result before applying the filter
Returns:
List of results.
* Implementation 2: DEPRECATED
Args:
values: Searched items (list[str])
limit: Upper limit for each result before applying the filter
Returns:
Dict with the searched terms as keys each with a list of results.
"""
return self._autosuggest_and_filter(
allowed_category=FilterAnalyticCategory.TOPICS, values=values, limit=limit
)
def _autosuggest_and_filter(
self,
allowed_category: FilterAnalyticCategory,
values: Union[list[str], str],
limit: int,
) -> Union[dict, list]:
if isinstance(values, list):
# Decorator not displaying the msg
warnings.warn(
DEPRECATED_WARNING_AUTOSUGGEST, DeprecationWarning, stacklevel=3
)
api_response = self._autosuggest_async(values, limit)
results = self._include_only_models(
api_response, models=(allowed_category.to_expected_model(),)
)
return dict(results.root.items())
if isinstance(values, str):
return self._autosuggest(
value=values, limit=limit, categories=[allowed_category]
)
raise TypeError(f"Expected list or str for @values, found {type(values)}")
@staticmethod
def _exclude_models(
api_response: Union[AutosuggestResponse, list[KnowledgeGraphTypes]],
models: tuple,
) -> Union[AutosuggestResponse, list[KnowledgeGraphTypes]]:
"""It will exclude the models from the response."""
if isinstance(api_response, AutosuggestResponse):
filtered_response = {}
for key, key_results in api_response.root.items():
filtered_response[key] = list(
filter(
lambda result: not isinstance(result, models),
key_results,
)
)
return AutosuggestResponse(root=filtered_response)
if isinstance(api_response, list):
return list(
filter(
lambda result: not isinstance(result, models),
api_response,
)
)
raise TypeError(
f"Expected AutosuggestResponse or list for @api_response, found {type(api_response)}"
)
@staticmethod
def _include_only_models(
api_response: Union[AutosuggestResponse, list[KnowledgeGraphTypes]],
models: tuple,
) -> Union[AutosuggestResponse, list[KnowledgeGraphTypes]]:
"""It will include the models specified only"""
if isinstance(api_response, AutosuggestResponse):
filtered_response = {}
for key, key_results in api_response.root.items():
filtered_response[key] = list(
filter(
lambda result: isinstance(result, models),
key_results,
)
)
return AutosuggestResponse(root=filtered_response)
if isinstance(api_response, list):
return list(
filter(
lambda result: isinstance(result, models),
api_response,
)
)
raise TypeError(
f"Expected AutosuggestResponse or list for @api_response, found {type(api_response)}"
)
[docs]
def get_entities(self, ids: list[str], /) -> list[Optional[EntityTypes]]:
"""Retrieve a list of entities by their ids."""
return self._get_by_ids(ids, QueryType.ENTITY)
[docs]
def get_sources(self, ids: list[str], /) -> list[Optional[Source]]:
"""Retrieve a list of sources by its ids."""
return self._get_by_ids(ids, QueryType.SOURCE)
[docs]
def get_topics(self, ids: list[str], /) -> list[Optional[Topic]]:
"""Retrieve a list of topics by its ids."""
return self._get_by_ids(ids, QueryType.TOPIC)
[docs]
def get_languages(self, ids: list[str], /) -> list[Optional[Language]]:
"""Retrieve a list of languages by its ids."""
return self._get_by_ids(ids, QueryType.LANGUAGE)
def _get_by_ids(self, ids: list[str], query_type: QueryType) -> list:
api_response = self._api.by_ids(
ByIdsRequest.model_validate(
[{"key": id_, "queryType": query_type} for id_ in ids]
)
)
return [api_response.root.get(id_) for id_ in ids]