Start implementing shape check

Author: peanutfun

climada/util/forecast.py

@@ -19,9 +19,28 @@
 Define Forecast base class.
 """
 
+from typing import Any
+
 import numpy as np
 
 
+def check_attribute_shapes(obj_act: Any, attr_act: str, obj_exp: Any, attr_exp: str):
+    """Compare the shapes of attributes of two objects.
+
+    Raises
+    ------
+    ValueError
+        If the shapes do not match
+    """
+    shape_actual = getattr(obj_act, attr_act).shape
+    shape_expected = getattr(obj_exp, attr_exp).shape
+    if shape_actual != shape_expected:
+        raise ValueError(
+            f"Shape mismatch between {type(obj_act).__name__}.{attr_act} "
+            f"{shape_actual} and {type(obj_exp).__name__}.{attr_exp} {shape_expected}"
+        )
+
+
 class Forecast:
     """Mixin class for forecast data.
 

climada/util/test/test_forecast.py

@@ -22,8 +22,9 @@
 import numpy as np
 import numpy.testing as npt
 import pandas as pd
+import pytest
 
-from climada.util.forecast import Forecast
+from climada.util.forecast import Forecast, check_attribute_shapes
 
 
 def test_forecast_init():
@@ -50,3 +51,32 @@ def test_forecast_init():
     forecast = Forecast(lead_time=lead_times_seconds, member=[1, 2, 3])
     npt.assert_array_equal(forecast.lead_time, lead_times_seconds, strict=True)
     assert forecast.lead_time.dtype == np.dtype("timedelta64[ns]")
+
+
+class A:
+    foo = np.array([[0, 1], [1, 0]])
+
+
+class B:
+    bar = np.array([[1, 1], [1, 1]])
+
+
+class TestCheckCompareShapes:
+    @pytest.fixture
+    def a(self):
+        return A()
+
+    @pytest.fixture
+    def b(self):
+        return B()
+
+    def test_pass(self, a, b):
+        check_attribute_shapes(a, "foo", b, "bar")
+
+    def test_error(self, a, b):
+        a.foo = np.array([0, 1])
+        with pytest.raises(ValueError, match=r"A.foo \(2\,\)"):
+            check_attribute_shapes(a, "foo", b, "bar")
+        b.bar = np.array([0, 1, 2])
+        with pytest.raises(ValueError, match=r"B.bar \(3\,\)"):
+            check_attribute_shapes(a, "foo", b, "bar")