kensho-kfinance 2.2.4__tar.gz → 2.2.5__tar.gz

{kensho_kfinance-2.2.4/kensho_kfinance.egg-info → kensho_kfinance-2.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kensho-kfinance
-Version: 2.2.4
+Version: 2.2.5
 Summary: Python CLI for kFinance
 Author-email: Luke Brown <luke.brown@kensho.com>, Michelle Keoy <michelle.keoy@kensho.com>, Keith Page <keith.page@kensho.com>, Matthew Rosen <matthew.rosen@kensho.com>, Nick Roshdieh <nick.roshdieh@kensho.com>
 Project-URL: source, https://github.com/kensho-technologies/kfinance

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5/kensho_kfinance.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kensho-kfinance
-Version: 2.2.4
+Version: 2.2.5
 Summary: Python CLI for kFinance
 Author-email: Luke Brown <luke.brown@kensho.com>, Michelle Keoy <michelle.keoy@kensho.com>, Keith Page <keith.page@kensho.com>, Matthew Rosen <matthew.rosen@kensho.com>, Nick Roshdieh <nick.roshdieh@kensho.com>
 Project-URL: source, https://github.com/kensho-technologies/kfinance

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kensho_kfinance.egg-info/SOURCES.txt

@@ -36,6 +36,7 @@ kfinance/kfinance.py
 kfinance/meta_classes.py
 kfinance/prompt.py
 kfinance/py.typed
+kfinance/pydantic_models.py
 kfinance/server_thread.py
 kfinance/version.py
 kfinance/tests/__init__.py

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kfinance/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v2.2.5
+- Add parsing for relationship response with name
+
 ## v2.2.4
 - Add segments as llm tools
 

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kfinance/fetch.py

@@ -6,6 +6,7 @@ from typing import Callable, Generator, Optional
 from uuid import uuid4
 
 import jwt
+from pydantic import ValidationError
 import requests
 
 from .constants import (
@@ -17,6 +18,7 @@ from .constants import (
     Permission,
     SegmentType,
 )
+from .pydantic_models import RelationshipResponse, RelationshipResponseNoName
 
 
 # version.py gets autogenerated by setuptools-scm and is not available
@@ -505,26 +507,30 @@ class KFinanceApiClient:
 
     def fetch_companies_from_business_relationship(
         self, company_id: int, relationship_type: BusinessRelationshipType
-    ) -> dict[str, list[int]]:
-        """Fetches a dictionary of current and previous company IDs associated with a given company ID based on the specified relationship type.
-
-        The returned dictionary has the following structure:
-        {
-            "current": List[int],
-            "previous": List[int]
-        }
+    ) -> RelationshipResponse | RelationshipResponseNoName:
+        """Fetches a dictionary of current and previous company IDs and names associated with a given company ID based on the specified relationship type.
 
         Example: fetch_companies_from_business_relationship(company_id=1234, relationship_type="distributor") returns a dictionary of company 1234's current and previous distributors.
 
+        As of 2024-05-28, we are changing the response on the backend from
+        RelationshipResponseNoName to RelationshipResponse. This function can handle both response
+        types.
+
         :param company_id: The ID of the company for which associated companies are being fetched.
         :type company_id: int
         :param relationship_type: The type of relationship to filter by. Valid relationship types are defined in the BusinessRelationshipType class.
         :type relationship_type: BusinessRelationshipType
         :return: A dictionary containing lists of current and previous company IDs that have the specified relationship with the given company_id.
-        :rtype: dict[str, list[int]]
+        :rtype: RelationshipResponse | RelationshipResponseNoName
         """
         url = f"{self.url_base}relationship/{company_id}/{relationship_type}"
-        return self.fetch(url)
+        result = self.fetch(url)
+        # Try to parse as the newer RelationshipResponse and fall back to
+        # RelationshipResponseNoName if that fails.
+        try:
+            return RelationshipResponse.model_validate(result)
+        except ValidationError:
+            return RelationshipResponseNoName.model_validate(result)
 
     def fetch_ticker_from_industry_code(
         self,

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kfinance/meta_classes.py

@@ -9,12 +9,12 @@ import pandas as pd
 
 from .constants import LINE_ITEMS, BusinessRelationshipType, PeriodType, SegmentType
 from .fetch import KFinanceApiClient
+from .pydantic_models import RelationshipResponse
 
 
 if TYPE_CHECKING:
     from .kfinance import BusinessRelationships
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -223,14 +223,32 @@ class CompanyFunctionsMetaClass:
         """
         from .kfinance import BusinessRelationships, Companies
 
-        companies = self.kfinance_api_client.fetch_companies_from_business_relationship(
-            self.company_id,
-            relationship_type,
-        )
-        return BusinessRelationships(
-            Companies(self.kfinance_api_client, companies["current"]),
-            Companies(self.kfinance_api_client, companies["previous"]),
+        relationship_resp = self.kfinance_api_client.fetch_companies_from_business_relationship(
+            company_id=self.company_id,
+            relationship_type=relationship_type,
         )
+        if isinstance(relationship_resp, RelationshipResponse):
+            return BusinessRelationships(
+                current=Companies(
+                    kfinance_api_client=self.kfinance_api_client,
+                    company_ids=[c.company_id for c in relationship_resp.current],
+                ),
+                previous=Companies(
+                    kfinance_api_client=self.kfinance_api_client,
+                    company_ids=[c.company_id for c in relationship_resp.previous],
+                ),
+            )
+        else:
+            return BusinessRelationships(
+                current=Companies(
+                    kfinance_api_client=self.kfinance_api_client,
+                    company_ids=relationship_resp.current,
+                ),
+                previous=Companies(
+                    kfinance_api_client=self.kfinance_api_client,
+                    company_ids=relationship_resp.previous,
+                ),
+            )
 
     def market_cap(
         self,

kensho_kfinance-2.2.5/kfinance/pydantic_models.py

@@ -0,0 +1,25 @@
+from pydantic import BaseModel
+
+
+class RelationshipResponseNoName(BaseModel):
+    """A response from the relationship endpoint before adding the company name.
+
+    Each element in `current` and `previous` is a company_id.
+    """
+
+    current: list[int]
+    previous: list[int]
+
+
+class CompanyIdAndName(BaseModel):
+    """A company_id and name"""
+
+    company_id: int
+    company_name: str
+
+
+class RelationshipResponse(BaseModel):
+    """A response from the relationship endpoint that includes both company_id and name."""
+
+    current: list[CompanyIdAndName]
+    previous: list[CompanyIdAndName]

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kfinance/tests/test_fetch.py

@@ -2,9 +2,17 @@ from unittest import TestCase
 from unittest.mock import MagicMock
 
 import pytest
+from requests_mock import Mocker
 
-from kfinance.constants import Periodicity, PeriodType
+from kfinance.constants import BusinessRelationshipType, Periodicity, PeriodType
 from kfinance.fetch import KFinanceApiClient
+from kfinance.kfinance import Client
+from kfinance.pydantic_models import (
+    CompanyIdAndName,
+    RelationshipResponse,
+    RelationshipResponseNoName,
+)
+from kfinance.tests.conftest import SPGI_COMPANY_ID
 
 
 def build_mock_api_client() -> KFinanceApiClient:
@@ -279,3 +287,56 @@ class TestMarketCap:
         expected_fetch_url = f"{client.url_base}users/permissions"
         client.fetch_permissions()
         client.fetch.assert_called_with(expected_fetch_url)
+
+
+class TestFetchCompaniesFromBusinessRelationship:
+    def test_old_response_format(self, requests_mock: Mocker, mock_client: Client) -> None:
+        """
+        GIVEN a business relationship request
+        WHEN the api returns a response in the old (no name) format
+        THEN the response can successfully be parsed.
+        """
+        http_resp = {"current": [883103], "previous": [472898, 8182358]}
+        expected_result = RelationshipResponseNoName(current=[883103], previous=[472898, 8182358])
+        requests_mock.get(
+            url=f"{mock_client.kfinance_api_client.url_base}relationship/{SPGI_COMPANY_ID}/{BusinessRelationshipType.supplier}",
+            json=http_resp,
+        )
+
+        resp = mock_client.kfinance_api_client.fetch_companies_from_business_relationship(
+            company_id=SPGI_COMPANY_ID, relationship_type=BusinessRelationshipType.supplier
+        )
+        assert resp == expected_result
+
+    def test_new_response_format(self, requests_mock: Mocker, mock_client: Client) -> None:
+        """
+        GIVEN a business relationship request
+        WHEN the api returns a response in the new (with name) format
+        THEN the response can successfully be parsed.
+        """
+
+        http_resp = {
+            "current": [{"company_name": "foo", "company_id": 883103}],
+            "previous": [
+                {"company_name": "bar", "company_id": 472898},
+                {"company_name": "baz", "company_id": 8182358},
+            ],
+        }
+
+        expected_result = RelationshipResponse(
+            current=[CompanyIdAndName(company_name="foo", company_id=883103)],
+            previous=[
+                CompanyIdAndName(company_name="bar", company_id=472898),
+                CompanyIdAndName(company_name="baz", company_id=8182358),
+            ],
+        )
+
+        requests_mock.get(
+            url=f"{mock_client.kfinance_api_client.url_base}relationship/{SPGI_COMPANY_ID}/{BusinessRelationshipType.supplier}",
+            json=http_resp,
+        )
+
+        resp = mock_client.kfinance_api_client.fetch_companies_from_business_relationship(
+            company_id=SPGI_COMPANY_ID, relationship_type=BusinessRelationshipType.supplier
+        )
+        assert resp == expected_result

{kensho_kfinance-2.2.4 → kensho_kfinance-2.2.5}/kfinance/version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.2.4'
-__version_tuple__ = version_tuple = (2, 2, 4)
+__version__ = version = '2.2.5'
+__version_tuple__ = version_tuple = (2, 2, 5)