fix other error handling

Author: matthew-rosen-12

kfinance/CHANGELOG.md

@@ -1,7 +1,7 @@
 # Changelog
 
 ## v5.1.1
-- Update estimates tools to gracefully handle errors
+- Update estimates, line item, statements, and segments tools to more gracefully handle errors
 
 ## v5.1.0
 - Add target price and analyst recommendation tools

kfinance/client/tests/test_fetch.py

@@ -425,7 +425,7 @@ def test_fetch_estimate(self) -> None:
         )
         assert isinstance(result, EstimatesResp)
         assert result.result is None
-        assert result.errors == {}
+        assert result.error is None
 
     def test_fetch_consensus_target_price(self) -> None:
         company_id = 21719
@@ -456,7 +456,6 @@ def test_fetch_consensus_target_price(self) -> None:
                     ),
                 ],
             ),
-            errors={},
         )
 
         result = self.kfinance_api_client.fetch_consensus_target_price(
@@ -492,7 +491,6 @@ def test_fetch_analyst_recommendations(self) -> None:
                     ),
                 ],
             ),
-            errors={},
         )
 
         result = self.kfinance_api_client.fetch_analyst_recommendations(

kfinance/client/tests/test_objects.py

@@ -499,17 +499,17 @@ def fetch_estimates(
         period_type,
     ):
         estimates = MOCK_COMPANY_DB[company_id]["estimates"]
-        return EstimatesResp(result=estimates, errors={})
+        return EstimatesResp(result=estimates)
 
     def fetch_consensus_target_price(self, company_id):
         """Get consensus target price estimates"""
         consensus_target_price = MOCK_COMPANY_DB[company_id]["consensus_target_price"]
-        return ConsensusTargetPriceResp(result=consensus_target_price, errors={})
+        return ConsensusTargetPriceResp(result=consensus_target_price)
 
     def fetch_analyst_recommendations(self, company_id):
         """Get analyst recommendations"""
         analyst_recommendations = MOCK_COMPANY_DB[company_id]["analyst_recommendations"]
-        return AnalystRecommendationsResp(result=analyst_recommendations, errors={})
+        return AnalystRecommendationsResp(result=analyst_recommendations)
 
     def fetch_line_item(
         self,

kfinance/domains/estimates/estimates_models.py

@@ -1,12 +1,11 @@
 from datetime import date
 from decimal import Decimal
 import logging
-from typing import Any
+from typing import Any, Callable, Dict
 
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, model_serializer, model_validator
 
 from kfinance.client.models.date_and_period_models import EstimatePeriodType, EstimateType
-from kfinance.client.models.response_models import RespWithErrors
 
 
 logger = logging.getLogger(__name__)
@@ -29,10 +28,20 @@ class Estimates(BaseModel):
     periods: dict[str, EstimatesPeriodData]
 
 
-class EstimatesResp(RespWithErrors):
+class EstimatesResp(BaseModel):
     """Response model for a single company's estimates."""
 
     result: Estimates | None = None
+    error: str | None = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler: Callable) -> Dict[str, Any]:
+        """Make `error` the last response field and only include if there is at least one error."""
+        data = handler(self)
+        error = data.pop("error")
+        if error:
+            data["error"] = error
+        return data
 
     @model_validator(mode="before")
     @classmethod
@@ -43,7 +52,11 @@ def from_post_response(cls, data: Any) -> Any:
             if len(results) > 1:
                 logger.warning("Expected at most one result, got %d", len(results))
             result = next(iter(results.values()), None)
-            return {"result": result, "errors": data.get("errors", {})}
+            errors = data.get("errors", {})
+            if len(errors) > 1:
+                logger.warning("Expected at most one error, got %d", len(errors))
+            error = next(iter(errors.values()), None)
+            return {"result": result, "error": error}
         return data
 
 
@@ -58,10 +71,20 @@ class ConsensusTargetPrice(BaseModel):
     estimates: list[ConsensusTargetPriceItem]
 
 
-class ConsensusTargetPriceResp(RespWithErrors):
+class ConsensusTargetPriceResp(BaseModel):
     """Response model for a single company's consensus target price."""
 
     result: ConsensusTargetPrice | None = None
+    error: str | None = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler: Callable) -> Dict[str, Any]:
+        """Make `error` the last response field and only include if there is at least one error."""
+        data = handler(self)
+        error = data.pop("error")
+        if error:
+            data["error"] = error
+        return data
 
     @model_validator(mode="before")
     @classmethod
@@ -72,7 +95,11 @@ def from_post_response(cls, data: Any) -> Any:
             if len(results) > 1:
                 logger.warning("Expected at most one result, got %d", len(results))
             result = next(iter(results.values()), None)
-            return {"result": result, "errors": data.get("errors", {})}
+            errors = data.get("errors", {})
+            if len(errors) > 1:
+                logger.warning("Expected at most one error, got %d", len(errors))
+            error = next(iter(errors.values()), None)
+            return {"result": result, "error": error}
         return data
 
 
@@ -86,10 +113,20 @@ class AnalystRecommendations(BaseModel):
     estimates: list[AnalystRecommendationsItem]
 
 
-class AnalystRecommendationsResp(RespWithErrors):
+class AnalystRecommendationsResp(BaseModel):
     """Response model for a single company's analyst recommendations."""
 
     result: AnalystRecommendations | None = None
+    error: str | None = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler: Callable) -> Dict[str, Any]:
+        """Make `error` the last response field and only include if there is at least one error."""
+        data = handler(self)
+        error = data.pop("error")
+        if error:
+            data["error"] = error
+        return data
 
     @model_validator(mode="before")
     @classmethod
@@ -100,5 +137,9 @@ def from_post_response(cls, data: Any) -> Any:
             if len(results) > 1:
                 logger.warning("Expected at most one result, got %d", len(results))
             result = next(iter(results.values()), None)
-            return {"result": result, "errors": data.get("errors", {})}
+            errors = data.get("errors", {})
+            if len(errors) > 1:
+                logger.warning("Expected at most one error, got %d", len(errors))
+            error = next(iter(errors.values()), None)
+            return {"result": result, "error": error}
         return data

kfinance/domains/estimates/estimates_tools.py

@@ -226,7 +226,9 @@ async def get_estimates_from_identifiers(
             resp: EstimatesResp = task.result
             if resp.result is not None:
                 results[task.result_key] = resp.result
-            errors.extend(resp.errors.values())
+            if resp.error is not None:
+                error_msg = f"{task.result_key}: {resp.error}"
+                errors.append(error_msg)
 
     resp_model = GetEstimatesFromIdentifiersResp(results=results, errors=errors)
 
@@ -311,7 +313,9 @@ async def get_consensus_target_price_from_identifiers(
             resp: ConsensusTargetPriceResp = task.result
             if resp.result is not None:
                 results[task.result_key] = resp.result
-            errors.extend(resp.errors.values())
+            if resp.error is not None:
+                error_msg = f"{task.result_key}: {resp.error}"
+                errors.append(error_msg)
 
     return GetConsensusTargetPriceFromIdentifiersResp(results=results, errors=errors)
 
@@ -358,7 +362,9 @@ async def get_analyst_recommendations_from_identifiers(
             resp: AnalystRecommendationsResp = task.result
             if resp.result is not None:
                 results[task.result_key] = resp.result
-            errors.extend(resp.errors.values())
+            if resp.error is not None:
+                error_msg = f"{task.result_key}: {resp.error}"
+                errors.append(error_msg)
 
     return GetAnalystRecommendationsFromIdentifiersResp(results=results, errors=errors)
 

kfinance/domains/estimates/tests/test_estimate_tools.py

@@ -98,7 +98,6 @@ async def test_fetch_estimates_from_company_id(
 
         expected_resp = EstimatesResp(
             result=Estimates.model_validate(self.estimates_data),
-            errors={},
         )
         assert resp == expected_resp
 
@@ -158,7 +157,6 @@ async def test_get_estimates_with_guidance_type(
 
         expected_resp = EstimatesResp(
             result=Estimates.model_validate(guidance_data),
-            errors={},
         )
         assert resp == expected_resp
         assert resp.result.estimate_type == EstimateType.guidance
@@ -200,7 +198,6 @@ async def test_fetch_consensus_target_price_from_company_id(
 
         expected_resp = ConsensusTargetPriceResp(
             result=ConsensusTargetPrice.model_validate(consensus_target_price_data),
-            errors={},
         )
         assert resp == expected_resp
 
@@ -282,7 +279,6 @@ async def test_fetch_analyst_recommendations_from_company_id(
 
         expected_resp = AnalystRecommendationsResp(
             result=AnalystRecommendations.model_validate(analyst_recommendations_data),
-            errors={},
         )
         assert resp == expected_resp
 

kfinance/domains/line_items/line_item_tools.py

@@ -5,9 +5,9 @@
 import httpx
 from pydantic import BaseModel, Field, model_validator
 
-from kfinance.async_batch_execution import AsyncTask, batch_execute_async_tasks
 from kfinance.client.id_resolution import unified_fetch_id_triples
 from kfinance.client.models.date_and_period_models import NumPeriods, NumPeriodsBack, PeriodType
+from kfinance.client.models.response_models import PostResponse
 from kfinance.client.permission_models import Permission
 from kfinance.domains.line_items.line_item_models import (
     LINE_ITEM_NAMES_AND_ALIASES,
@@ -241,40 +241,38 @@ async def get_financial_line_item_from_identifiers(
         id_triple.company_id for id_triple in id_triple_resp.identifiers_to_id_triples.values()
     ]
 
-    # Create a single task to fetch line items for all company IDs at once
+    # Fetch line items for all company IDs at once
     if company_ids:
-        task = AsyncTask(
-            func=fetch_line_item_from_company_ids,
-            kwargs=dict(
-                company_ids=company_ids,
-                line_item=line_item,
-                httpx_client=httpx_client,
-                period_type=period_type,
-                start_year=start_year,
-                end_year=end_year,
-                start_quarter=start_quarter,
-                end_quarter=end_quarter,
-                calendar_type=calendar_type,
-                num_periods=num_periods,
-                num_periods_back=num_periods_back,
-            ),
-            result_key="line_items",
+        line_item_resp = await fetch_line_item_from_company_ids(
+            company_ids=company_ids,
+            line_item=line_item,
+            httpx_client=httpx_client,
+            period_type=period_type,
+            start_year=start_year,
+            end_year=end_year,
+            start_quarter=start_quarter,
+            end_quarter=end_quarter,
+            calendar_type=calendar_type,
+            num_periods=num_periods,
+            num_periods_back=num_periods_back,
         )
 
-        await batch_execute_async_tasks(tasks=[task])
-
-        if task.error:
-            errors.append(task.error)
-            results = {}
-        else:
-            # Map company IDs back to original identifiers
-            identifier_to_results = {}
-            for company_id_str, line_item_resp in task.result.items():
-                original_identifier = id_triple_resp.get_identifier_from_company_id(
-                    int(company_id_str)
-                )
-                identifier_to_results[original_identifier] = line_item_resp
-            results = identifier_to_results
+        # Add any errors from the line item API, mapping company_id keys back to identifiers
+        for company_id_str, error in line_item_resp.errors.items():
+            original_identifier = id_triple_resp.get_identifier_from_company_id(
+                int(company_id_str)
+            )
+            errors.append(f"{original_identifier}: {error}")
+
+        # Map results back to original identifiers
+        identifier_to_results = {}
+        for company_id_str, line_item_data in line_item_resp.results.items():
+            original_identifier = id_triple_resp.get_identifier_from_company_id(
+                int(company_id_str)
+            )
+            identifier_to_results[original_identifier] = line_item_data
+
+        results = identifier_to_results
     else:
         results = {}
 
@@ -316,7 +314,7 @@ async def fetch_line_item_from_company_ids(
     calendar_type: CalendarType | None = None,
     num_periods: int | None = None,
     num_periods_back: int | None = None,
-) -> dict[str, LineItemResp]:
+) -> PostResponse[LineItemResp]:
     """Fetch line items for a list of company IDs."""
     # Build the request payload
     params: dict[str, Any] = {
@@ -342,11 +340,5 @@ async def fetch_line_item_from_company_ids(
         params["num_periods_back"] = num_periods_back
 
     resp = await httpx_client.post(url="/line_item/", json=params)
-    response_data = resp.json()
-
-    # Convert the response data to LineItemResp objects
-    results = {}
-    for company_id_str, line_item_data in response_data["results"].items():
-        results[company_id_str] = LineItemResp.model_validate(line_item_data)
 
-    return results
+    return PostResponse[LineItemResp].model_validate(resp.json())

kfinance/domains/line_items/tests/test_line_item_tools.py

@@ -70,10 +70,11 @@ async def test_fetch_line_item_from_company_ids(
             httpx_client=httpx_client,
         )
 
-        expected_resp = {
+        expected_results = {
             str(SPGI_ID_TRIPLE.company_id): LineItemResp.model_validate(self.line_item_resp)
         }
-        assert resp == expected_resp
+        assert resp.results == expected_results
+        assert resp.errors == {}
 
     @pytest.mark.parametrize(
         "calendar_type, expected_notes",

kfinance/domains/segments/segment_tools.py

@@ -136,8 +136,10 @@ async def get_segments_from_identifiers(
         httpx_client=httpx_client,
     )
 
-    # Add any errors from the segments API
-    errors.extend(segments_resp.errors.values())
+    # Add any errors from the segments API, mapping company_id keys back to identifiers
+    for company_id_str, error in segments_resp.errors.items():
+        original_identifier = id_triple_resp.get_identifier_from_company_id(int(company_id_str))
+        errors.append(f"{original_identifier}: {error}")
 
     # Map results back to original identifiers
     identifier_to_results = {}

kfinance/domains/statements/statement_tools.py

@@ -160,8 +160,10 @@ async def get_financial_statement_from_identifiers(
         httpx_client=httpx_client,
     )
 
-    # Add any errors from the statements API
-    errors.extend(statements_resp.errors.values())
+    # Add any errors from the statements API, mapping company_id keys back to identifiers
+    for company_id_str, error in statements_resp.errors.items():
+        original_identifier = id_triple_resp.get_identifier_from_company_id(int(company_id_str))
+        errors.append(f"{original_identifier}: {error}")
 
     # Map results back to original identifiers
     identifier_to_results = {}