feat: allow state method passthrough in MockScenario objects

powersimdata/tests/mock_analyze.py

@@ -23,6 +23,39 @@ def _ensure_ts_index(df):
 
 
 class MockAnalyze:
+    """A mock of a powersimdata.scenario.analyze.Analyze object.
+
+    :param dict grid_attrs: fields to be added to grid.
+    :param pandas.DataFrame congl: dummy congl
+    :param pandas.DataFrame congu: dummy congu
+    :param dict ct: dummy ct
+    :param pandas.DataFrame demand: dummy demand
+    :param pandas.DataFrame lmp: dummy lmp
+    :param pandas.DataFrame pf: dummy pf
+    :param pandas.DataFrame pg: dummy pg
+    :param pandas.DataFrame storage_pg: dummy storage_pg
+    :param pandas.DataFrame solar: dummy solar
+    :param pandas.DataFrame wind: dummy wind
+    :param pandas.DataFrame hydro: dummy hydro
+    """
+
+    exported_methods = [
+        "get_congl",
+        "get_congu",
+        "get_ct",
+        "get_demand",
+        "get_grid",
+        "get_lmp",
+        "get_pf",
+        "gt_dcline_pf",
+        "get_pg",
+        "get_storage_pg",
+        "get_storage_e",
+        "get_solar",
+        "get_wind",
+        "get_hydro",
+    ]
+
     def __init__(
         self,
         grid_attrs=None,
@@ -34,26 +67,13 @@ def __init__(
         pf=None,
         pg=None,
         dcline_pf=None,
+        storage_e=None,
         storage_pg=None,
         solar=None,
         wind=None,
         hydro=None,
     ):
-        """Constructor.
-
-        :param dict grid_attrs: fields to be added to grid.
-        :param pandas.DataFrame congl: dummy congl
-        :param pandas.DataFrame congu: dummy congu
-        :param dict ct: dummy ct
-        :param pandas.DataFrame demand: dummy demand
-        :param pandas.DataFrame lmp: dummy lmp
-        :param pandas.DataFrame pf: dummy pf
-        :param pandas.DataFrame pg: dummy pg
-        :param pandas.DataFrame storage_pg: dummy storage_pg
-        :param pandas.DataFrame solar: dummy solar
-        :param pandas.DataFrame wind: dummy wind
-        :param pandas.DataFrame hydro: dummy hydro
-        """
+        """Constructor."""
         self.grid = MockGrid(grid_attrs)
         self.congl = _ensure_ts_index(congl)
         self.congu = _ensure_ts_index(congu)
@@ -63,6 +83,7 @@ def __init__(
         self.pf = _ensure_ts_index(pf)
         self.dcline_pf = _ensure_ts_index(dcline_pf)
         self.pg = _ensure_ts_index(pg)
+        self.storage_e = _ensure_ts_index(storage_e)
         self.storage_pg = _ensure_ts_index(storage_pg)
         self.solar = _ensure_ts_index(solar)
         self.wind = _ensure_ts_index(wind)
@@ -123,6 +144,12 @@ def get_pg(self):
         """
         return self.pg
 
+    def get_storage_e(self):
+        """Get storage E.
+        :return: (pandas.DataFrame) -- dummy storage_e
+        """
+        return self.storage_e
+
     def get_storage_pg(self):
         """Get storage PG.
         :return: (pandas.DataFrame) -- dummy storage_pg

powersimdata/tests/mock_scenario.py

@@ -3,12 +3,16 @@
 
 
 class MockScenario:
-    def __init__(self, grid_attrs=None, **kwargs):
-        """Constructor.
+    """
+    :param dict grid_attrs: fields to be added to grid.
+    :param \\*\\*kwargs: collected keyword arguments to be passed to
+        MockAnalyze init.
+    """
+
+    _setattr_allowlist = {"state", "info"}
 
-        :param dict grid_attrs: fields to be added to grid.
-        :param pandas.DataFrame pg: dummy pg
-        """
+    def __init__(self, grid_attrs=None, **kwargs):
+        """Constructor."""
         self.state = MockAnalyze(grid_attrs, **kwargs)
         self.info = {
             "id": "111",
@@ -34,3 +38,26 @@ def __init__(self, grid_attrs=None, **kwargs):
     def __class__(self):
         """If anyone asks, I'm a Scenario object!"""
         return Scenario
+
+    def __dir__(self):
+        return sorted(super().__dir__() + list(self.state.exported_methods))
+
+    def __getattr__(self, name):
+        if name in self.state.exported_methods:
+            return getattr(self.state, name)
+        elif hasattr(self.state, "__getattr__"):
+            return self.state.__getattr__(name)
+        else:
+            raise AttributeError(
+                f"Scenario object in {self.state.name} state "
+                f"has no attribute {name}"
+            )
+
+    def __setattr__(self, name, value):
+        if name in self._setattr_allowlist:
+            super().__setattr__(name, value)
+        elif name in self.state.exported_methods:
+            raise AttributeError(
+                f"{name} is exported from Scenario.state, edit it there if necessary"
+            )
+        super().__setattr__(name, value)