From 882cefc0533f096e96d796f20bda02dfea4102b0 Mon Sep 17 00:00:00 2001
From: Jen Hagg <jen@breakthroughenergy.org>
Date: Fri, 2 Dec 2022 16:36:35 -0800
Subject: [PATCH] refactor: add back constructor to FromPyPSA

---
 powersimdata/input/converter/pypsa_to_grid.py          |  6 ++++++
 .../input/converter/tests/test_pypsa_to_grid.py        | 10 ++--------
 powersimdata/network/europe_tub/model.py               |  4 +---
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/powersimdata/input/converter/pypsa_to_grid.py b/powersimdata/input/converter/pypsa_to_grid.py
index c67389bd4..599444cff 100644
--- a/powersimdata/input/converter/pypsa_to_grid.py
+++ b/powersimdata/input/converter/pypsa_to_grid.py
@@ -131,6 +131,12 @@ def _get_storage_gen(df, storage_type):
 class FromPyPSA(AbstractGrid):
     """Grid builder for PyPSA network object."""
 
+    def __init__(self, network, add_pypsa_cols=True):
+        """Constructor."""
+        super().__init__()
+        self.network = network
+        self.add_pypsa_cols = add_pypsa_cols
+
     def _set_interconnect(self):
         if self.interconnect is None:
             self.interconnect = self.network.name.split(", ")
diff --git a/powersimdata/input/converter/tests/test_pypsa_to_grid.py b/powersimdata/input/converter/tests/test_pypsa_to_grid.py
index c962dec73..b505dadc5 100644
--- a/powersimdata/input/converter/tests/test_pypsa_to_grid.py
+++ b/powersimdata/input/converter/tests/test_pypsa_to_grid.py
@@ -11,10 +11,7 @@
 def test_import_arbitrary_network_from_pypsa_to_grid():
 
     n = pypsa.examples.ac_dc_meshed()
-    grid = FromPyPSA()
-    grid.network = n
-    grid.add_pypsa_cols = True
-    grid.build()
+    grid = FromPyPSA(n).build()
 
     assert not grid.bus.empty
     assert len(n.buses) == len(grid.bus)
@@ -23,10 +20,7 @@ def test_import_arbitrary_network_from_pypsa_to_grid():
 def test_import_network_including_storages_from_pypsa_to_grid():
 
     n = pypsa.examples.storage_hvdc()
-    grid = FromPyPSA()
-    grid.network = n
-    grid.add_pypsa_cols = True
-    grid.build()
+    grid = FromPyPSA(n).build()
 
     inflow = n.get_switchable_as_dense("StorageUnit", "inflow")
     has_inflow = inflow.any()
diff --git a/powersimdata/network/europe_tub/model.py b/powersimdata/network/europe_tub/model.py
index d03ad0faf..57ef96f50 100644
--- a/powersimdata/network/europe_tub/model.py
+++ b/powersimdata/network/europe_tub/model.py
@@ -23,13 +23,11 @@ class PyPSABase(FromPyPSA):
 
     def __init__(self, interconnect, grid_model, network, add_pypsa_cols=True):
         """Constructor."""
-        super().__init__()
+        super().__init__(network, add_pypsa_cols)
         self.grid_model = grid_model
         self.interconnect = check_and_format_interconnect(
             interconnect, model=self.grid_model
         )
-        self.network = network
-        self.add_pypsa_cols = add_pypsa_cols
 
     def build_eur(self):
         self.id2zone = {i: l for i, l in enumerate(self.network.buses.index)}