Update all models to use Warnings

Author: louispt1

src/pyetm/models/base.py

@@ -92,16 +92,9 @@ def add_warning(
     def warnings(self) -> Union[WarningCollector, Dict[str, List[str]]]:
         """
         Return warnings.
-
-        For backward compatibility, this can return either the new WarningCollector
-        or the legacy dict format. The implementation can be switched based on needs.
         """
-        # Return the new collector (recommended)
         return self._warning_collector
 
-        # OR return legacy format for backward compatibility:
-        # return self._warning_collector.to_legacy_dict()
-
     def show_warnings(self) -> None:
         """Print all warnings to the console."""
         self._warning_collector.show_warnings()

src/pyetm/models/custom_curves.py

@@ -4,6 +4,7 @@
 from typing import Optional
 from pyetm.clients import BaseClient
 from pyetm.models.base import Base
+from pyetm.models.warnings import WarningCollector
 from pyetm.services.scenario_runners.fetch_custom_curves import (
     DownloadCustomCurveRunner,
 )
@@ -59,10 +60,7 @@ def retrieve(self, client, scenario) -> Optional[pd.DataFrame]:
                     return curve.rename(self.key)
                 except Exception as e:
                     # File processing error - add warning and return None
-                    self.add_warning(
-                        self.key,
-                        f"Failed to process curve data: {e}"
-                    )
+                    self.add_warning(self.key, f"Failed to process curve data: {e}")
                     return None
             else:
                 # API call failed - add warning for each error
@@ -119,8 +117,8 @@ def from_json(cls, data: dict) -> CustomCurve:
                 "key": data.get("key", "unknown"),
                 "type": data.get("type", "unknown"),
             }
-            curve = cls.model_validate(basic_data)
-            curve.add_warning(basic_data["key"], f"Failed to create curve from data: {e}")
+            curve = cls.model_construct(**basic_data)
+            curve.add_warning("base", f"Failed to create curve from data: {e}")
             return curve
 
 
@@ -147,19 +145,17 @@ def get_contents(
         curve = self._find(curve_name)
 
         if not curve:
-            self.add_warning('curves', f"Curve {curve_name} not found in collection")
+            self.add_warning("curves", f"Curve {curve_name} not found in collection")
             return None
 
         if not curve.available():
             # Try to retrieve it
             result = curve.retrieve(BaseClient(), scenario)
-            # Merge any warnings from the curve retrieval
-            self._merge_submodel_warnings(curve)
+            self._merge_submodel_warnings(curve, key_attr="key")
             return result
         else:
             contents = curve.contents()
-            # Merge any warnings from reading contents
-            self._merge_submodel_warnings(curve)
+            self._merge_submodel_warnings(curve, key_attr="key")
             return contents
 
     def _find(self, curve_name: str) -> Optional[CustomCurve]:
@@ -171,24 +167,23 @@ def from_json(cls, data: list[dict]) -> CustomCurves:
         Initialize CustomCurves collection from JSON data
         """
         curves = []
-        collection_warnings = {}
 
         for curve_data in data:
             try:
-                key = curve_data['key']
                 curve = CustomCurve.from_json(curve_data)
                 curves.append(curve)
             except Exception as e:
-                # Log the problematic curve but continue processing
-                collection_warnings[f"CustomCurve(key={key})"] = f"Skipped invalid curve data: {e}"
+                # Create a basic curve and continue processing
+                key = curve_data.get("key", "unknown")
+                basic_curve = CustomCurve.model_construct(key=key, type="unknown")
+                basic_curve.add_warning(key, f"Skipped invalid curve data: {e}")
+                curves.append(basic_curve)
 
         collection = cls.model_validate({"curves": curves})
 
-        # Add any collection-level warnings
-        for loc, msg in collection_warnings.items():
-            collection.add_warning(loc, msg)
-
-        # Merge warnings from individual curves
-        collection._merge_submodel_warnings(*curves, key_attr='key')
+        # Merge warnings from individual curves using new system
+        collection._merge_submodel_warnings(*curves, key_attr="key")
 
         return collection
+
+    # TODO: _to_dataframe

src/pyetm/models/gqueries.py

@@ -10,7 +10,7 @@
 
 class Gqueries(Base):
     """
-    We cannot validat yet - as we'd need a servcie connected to the main
+    We cannot validate yet - as we'd need a service connected to the main
     gquery endpoint
     """
 
@@ -48,7 +48,7 @@ def execute(self, client, scenario):
         if result.success:
             self.update(result.data)
         else:
-            self.add_warning('results', f"Error retrieving queries: {result.errors}")
+            self.add_warning("results", f"Error retrieving queries: {result.errors}")
 
     def to_dataframe(self, columns="future"):
         if not self.is_ready():
@@ -58,6 +58,13 @@ def to_dataframe(self, columns="future"):
         df.index.name = "gquery"
         return df.set_index("unit", append=True)
 
+    def _to_dataframe(self, **kwargs) -> pd.DataFrame:
+        """
+        Implementation required by Base class.
+        Uses to_dataframe with default parameters.
+        """
+        return self.to_dataframe()
+
     @classmethod
     def from_list(cls, query_list: list[str]):
         return cls(query_dict={q: None for q in query_list})

src/pyetm/models/inputs.py

@@ -45,7 +45,6 @@ def from_json(cls, data: tuple[str, dict]) -> "Input":
         except Exception as e:
             # Create a basic Input with warning attached
             basic_input = cls.model_construct(**payload)  # Bypass validation
-            basic_input._warning_collector = WarningCollector()
             basic_input.add_warning(key, f"Failed to create specialized input: {e}")
             return basic_input
 
@@ -114,7 +113,7 @@ def check_permitted(self) -> EnumInput:
             "ValueError",
             type="inclusion",
             loc="user",
-            msg=f"Value error, {self.user} should be in {self.permitted_values}",
+            msg=f"Value error, {self.value} should be in {self.permitted_values}",
         )
 
 
@@ -180,10 +179,7 @@ def is_valid_update(self, key_vals: dict) -> dict[str, WarningCollector]:
         # Check for non-existent keys
         non_existent_keys = set(key_vals.keys()) - set(self.keys())
         for key in non_existent_keys:
-            # Create a warning collector for non-existent keys
-            warning_collector = WarningCollector()
-            warning_collector.add(key, "Key does not exist")
-            warnings[key] = warning_collector
+            warnings[key] = WarningCollector.with_warning(key, "Key does not exist")
 
         return warnings
 

src/pyetm/models/output_curves.py

@@ -7,6 +7,7 @@
 import yaml
 from pyetm.clients import BaseClient
 from pyetm.models.base import Base
+from pyetm.models.warnings import WarningCollector
 from pyetm.config.settings import get_settings
 from pyetm.services.scenario_runners.fetch_output_curves import (
     DownloadOutputCurveRunner,
@@ -59,27 +60,32 @@ def retrieve(self, client, scenario) -> Optional[pd.DataFrame]:
 
                 except Exception as e:
                     self.add_warning(
-                        'data',
-                        f"Failed to process curve data for {self.key}: {e}"
+                        "data", f"Failed to process curve data for {self.key}: {e}"
                     )
                     return None
 
         except Exception as e:
             # Unexpected error - add warning
-            self.add_warning('base', f"Unexpected error retrieving curve {self.key}: {e}")
+            self.add_warning(
+                "base", f"Unexpected error retrieving curve {self.key}: {e}"
+            )
             return None
 
     def contents(self) -> Optional[pd.DataFrame]:
         """Open file from path and return contents"""
         if not self.available():
-            self.add_warning('file_path', f"Curve {self.key} not available - no file path set")
+            self.add_warning(
+                "file_path", f"Curve {self.key} not available - no file path set"
+            )
             return None
 
         try:
             df = pd.read_csv(self.file_path, index_col=0)
             return df.dropna(how="all")
         except Exception as e:
-            self.add_warning('file_path', f"Failed to read curve file for {self.key}: {e}")
+            self.add_warning(
+                "file_path", f"Failed to read curve file for {self.key}: {e}"
+            )
             return None
 
     def remove(self) -> bool:
@@ -92,7 +98,9 @@ def remove(self) -> bool:
             self.file_path = None
             return True
         except Exception as e:
-            self.add_warning('file_path', f"Failed to remove curve file for {self.key}: {e}")
+            self.add_warning(
+                "file_path", f"Failed to remove curve file for {self.key}: {e}"
+            )
             return False
 
     @classmethod
@@ -109,8 +117,8 @@ def from_json(cls, data: dict) -> OutputCurve:
                 "key": data.get("key", "unknown"),
                 "type": data.get("type", "unknown"),
             }
-            curve = cls.model_validate(basic_data)
-            curve.add_warning('base', f"Failed to create curve from data: {e}")
+            curve = cls.model_construct(**basic_data)
+            curve.add_warning("base", f"Failed to create curve from data: {e}")
             return curve
 
 
@@ -135,19 +143,19 @@ def get_contents(self, scenario, curve_name: str) -> Optional[pd.DataFrame]:
         curve = self._find(curve_name)
 
         if not curve:
-            self.add_warning('curves', f"Curve {curve_name} not found in collection")
+            self.add_warning("curves", f"Curve {curve_name} not found in collection")
             return None
 
         if not curve.available():
             # Try to retrieve it
             result = curve.retrieve(BaseClient(), scenario)
-            # Merge any warnings from the curve retrieval
-            self._merge_submodel_warnings(curve)
+            # Merge any warnings from the curve retrieval using new system
+            self._merge_submodel_warnings(curve, key_attr="key")
             return result
         else:
             contents = curve.contents()
-            # Merge any warnings from reading contents
-            self._merge_submodel_warnings(curve)
+            # Merge any warnings from reading contents using new system
+            self._merge_submodel_warnings(curve, key_attr="key")
             return contents
 
     @staticmethod
@@ -203,8 +211,8 @@ def get_curves_by_carrier_type(
         if carrier_type not in carrier_mapping:
             valid_types = ", ".join(carrier_mapping.keys())
             self.add_warning(
-                'carrier_type',
-                f"Invalid carrier type '{carrier_type}'. Valid types: {valid_types}"
+                "carrier_type",
+                f"Invalid carrier type '{carrier_type}'. Valid types: {valid_types}",
             )
             return {}
 
@@ -225,25 +233,22 @@ def from_json(cls, data: list[dict]) -> OutputCurves:
         Initialize OutputCurves collection from JSON data
         """
         curves = []
-        collection_warnings = {}
 
         for curve_data in data:
             try:
-                key = curve_data['key']
                 curve = OutputCurve.from_json(curve_data)
                 curves.append(curve)
             except Exception as e:
-                # Log the problematic curve but continue processing
-                collection_warnings[f"OutputCurve(key={key})"] = f"Skipped invalid curve data: {e}"
+                # Create a basic curve and continue processing
+                key = curve_data.get("key", "unknown")
+                basic_curve = OutputCurve.model_construct(key=key, type="unknown")
+                basic_curve.add_warning(key, f"Skipped invalid curve data: {e}")
+                curves.append(basic_curve)
 
         collection = cls.model_validate({"curves": curves})
 
-        # Add any collection-level warnings
-        for loc, msg in collection_warnings.items():
-            collection.add_warning(loc, msg)
-
-        # Merge warnings from individual curves
-        collection._merge_submodel_warnings(*curves, key_attr='key')
+        # Merge warnings from individual curves using new system
+        collection._merge_submodel_warnings(*curves, key_attr="key")
 
         return collection
 
@@ -255,7 +260,7 @@ def from_service_result(
         if not service_result.success or not service_result.data:
             empty_curves = cls(curves=[])
             for error in service_result.errors:
-                empty_curves.add_warning('base', f"Service error: {error}")
+                empty_curves.add_warning("base", f"Service error: {error}")
             return empty_curves
 
         curves_list = []
@@ -281,17 +286,18 @@ def from_service_result(
                 curves_list.append(curve)
 
             except Exception as e:
-                curves_list.append(
-                    OutputCurve.model_validate({"key": curve_name, "type": "unknown"})
+                basic_curve = OutputCurve.model_construct(
+                    key=curve_name, type="unknown"
                 )
-                curves_list[-1].add_warning('base', f"Failed to process curve data: {e}")
+                basic_curve.add_warning("base", f"Failed to process curve data: {e}")
+                curves_list.append(basic_curve)
 
         curves_collection = cls(curves=curves_list)
 
         for error in service_result.errors:
-            curves_collection.add_warning('base', f"Download warning: {error}")
+            curves_collection.add_warning("base", f"Download warning: {error}")
 
-        curves_collection._merge_submodel_warnings(curves_list, key_attr='key')
+        curves_collection._merge_submodel_warnings(*curves_list, key_attr="key")
 
         return curves_collection