diff --git a/tests/test_nodes.py b/tests/test_nodes.py index fd7ab690..81f9e93d 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -663,61 +663,82 @@ def test_decay_queue(self): def test_overrides(self): # node - no need to test # tank - tank = Tank(capacity=10, area=8, datum = 4) - tank.apply_overrides({'capacity': 3, - 'area': 2, - 'datum': 3.5}) + tank = Tank(capacity=10, area=8, datum=4) + tank.apply_overrides({"capacity": 3, "area": 2, "datum": 3.5}) self.assertEqual(tank.capacity, 3) self.assertEqual(tank.area, 2) self.assertEqual(tank.datum, 3.5) # residence tank - tank = ResidenceTank(capacity=10, area=8, datum = 4, residence_time = 8) - tank.apply_overrides({'capacity': 3, - 'area': 2, - 'datum': 3.5, - 'residence_time': 6}) + tank = ResidenceTank(capacity=10, area=8, datum=4, residence_time=8) + tank.apply_overrides( + {"capacity": 3, "area": 2, "datum": 3.5, "residence_time": 6} + ) self.assertEqual(tank.capacity, 3) self.assertEqual(tank.area, 2) self.assertEqual(tank.datum, 3.5) self.assertEqual(tank.residence_time, 6) # decay tank - tank = DecayTank(capacity=10, area=8, datum = 4, decays = {'nitrate' : {'constant' : 0.001, 'exponent' : 1.005}}) - tank.apply_overrides({'capacity': 3, - 'area': 2, - 'datum': 3.5, - 'decays': {'phosphate' : {'constant' : 1.001, 'exponent' : 10.005}} - }) + tank = DecayTank( + capacity=10, + area=8, + datum=4, + decays={"nitrate": {"constant": 0.001, "exponent": 1.005}}, + ) + tank.apply_overrides( + { + "capacity": 3, + "area": 2, + "datum": 3.5, + "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + } + ) self.assertEqual(tank.capacity, 3) self.assertEqual(tank.area, 2) self.assertEqual(tank.datum, 3.5) - self.assertDictEqual(tank.decays, {'nitrate' : {'constant' : 0.001, 'exponent' : 1.005}, - 'phosphate' : {'constant' : 1.001, 'exponent' : 10.005} - }) + self.assertDictEqual( + tank.decays, + { + "nitrate": {"constant": 0.001, "exponent": 1.005}, + "phosphate": {"constant": 1.001, "exponent": 10.005}, + }, + ) # queue tank - tank = QueueTank(capacity=10, area=8, datum = 4, number_of_timesteps = 8) - tank.apply_overrides({'capacity': 3, - 'area': 2, - 'datum': 3.5, - 'number_of_timesteps': 6}) + tank = QueueTank(capacity=10, area=8, datum=4, number_of_timesteps=8) + tank.apply_overrides( + {"capacity": 3, "area": 2, "datum": 3.5, "number_of_timesteps": 6} + ) self.assertEqual(tank.capacity, 3) self.assertEqual(tank.area, 2) self.assertEqual(tank.datum, 3.5) self.assertEqual(tank.number_of_timesteps, 6) self.assertEqual(tank.internal_arc.number_of_timesteps, 6) # decay queue tank - tank = DecayQueueTank(capacity=10, area=8, datum = 4, number_of_timesteps = 8, decays = {'phosphate' : {'constant' : 0.001, 'exponent' : 1.005}}) - tank.apply_overrides({'capacity': 3, - 'area': 2, - 'datum': 3.5, - 'number_of_timesteps': 6, - 'decays': {'phosphate' : {'constant' : 1.001, 'exponent' : 10.005}} - }) + tank = DecayQueueTank( + capacity=10, + area=8, + datum=4, + number_of_timesteps=8, + decays={"phosphate": {"constant": 0.001, "exponent": 1.005}}, + ) + tank.apply_overrides( + { + "capacity": 3, + "area": 2, + "datum": 3.5, + "number_of_timesteps": 6, + "decays": {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + } + ) self.assertEqual(tank.capacity, 3) self.assertEqual(tank.area, 2) self.assertEqual(tank.datum, 3.5) self.assertEqual(tank.number_of_timesteps, 6) self.assertEqual(tank.internal_arc.number_of_timesteps, 6) - self.assertDictEqual(tank.internal_arc.decays, {'phosphate' : {'constant' : 1.001, 'exponent' : 10.005}}) + self.assertDictEqual( + tank.internal_arc.decays, + {"phosphate": {"constant": 1.001, "exponent": 10.005}}, + ) + if __name__ == "__main__": unittest.main() diff --git a/wsimod/nodes/nodes.py b/wsimod/nodes/nodes.py index d01195ad..6606add3 100644 --- a/wsimod/nodes/nodes.py +++ b/wsimod/nodes/nodes.py @@ -793,25 +793,22 @@ def __init__(self, capacity=0, area=1, datum=10, initial_storage=0): else: self.storage = self.empty_vqip() self.storage_ = self.empty_vqip() # Lagged storage for mass balance - + def apply_overrides(self, overrides: Dict[str, Any] = {}): """Apply overrides to the tank. - Enables a user to override any of the following parameters: + Enables a user to override any of the following parameters: area, capacity, datum. - + Args: overrides (dict, optional): Dictionary of overrides. Defaults to {}. """ - self.capacity = overrides.pop("capacity", - self.capacity) - self.area = overrides.pop("area", - self.area) - self.datum = overrides.pop("datum", - self.datum) + self.capacity = overrides.pop("capacity", self.capacity) + self.area = overrides.pop("area", self.area) + self.datum = overrides.pop("datum", self.datum) if len(overrides) > 0: print(f"No override behaviour defined for: {overrides.keys()}") - + def ds(self): """Should be called by parent object to get change in storage. @@ -1119,15 +1116,14 @@ def __init__(self, residence_time=2, **kwargs): def apply_overrides(self, overrides: Dict[str, Any] = {}): """Apply overrides to the residencetank. - - Enables a user to override any of the following parameters: + + Enables a user to override any of the following parameters: residence_time. - + Args: overrides (dict, optional): Dictionary of overrides. Defaults to {}. """ - self.residence_time = overrides.pop("residence_time", - self.residence_time) + self.residence_time = overrides.pop("residence_time", self.residence_time) super().apply_overrides(overrides) def pull_outflow(self): @@ -1175,10 +1171,10 @@ def __init__(self, decays={}, parent=None, **kwargs): def apply_overrides(self, overrides: Dict[str, Any] = {}): """Apply overrides to the decaytank. - - Enables a user to override any of the following parameters: + + Enables a user to override any of the following parameters: decays. - + Args: overrides (dict, optional): Dictionary of overrides. Defaults to {}. """ @@ -1235,15 +1231,16 @@ def __init__(self, number_of_timesteps=0, **kwargs): def apply_overrides(self, overrides: Dict[str, Any] = {}): """Apply overrides to the queuetank. - - Enables a user to override any of the following parameters: + + Enables a user to override any of the following parameters: number_of_timesteps. - + Args: overrides (dict, optional): Dictionary of overrides. Defaults to {}. """ - self.number_of_timesteps = overrides.pop("number_of_timesteps", - self.number_of_timesteps) + self.number_of_timesteps = overrides.pop( + "number_of_timesteps", self.number_of_timesteps + ) self.internal_arc.number_of_timesteps = self.number_of_timesteps super().apply_overrides(overrides) @@ -1417,15 +1414,16 @@ def __init__(self, decays={}, parent=None, number_of_timesteps=1, **kwargs): def apply_overrides(self, overrides: Dict[str, Any] = {}): """Apply overrides to the decayqueuetank. - - Enables a user to override any of the following parameters: + + Enables a user to override any of the following parameters: number_of_timesteps, decays. - + Args: overrides (dict, optional): Dictionary of overrides. Defaults to {}. """ - self.number_of_timesteps = overrides.pop("number_of_timesteps", - self.number_of_timesteps) + self.number_of_timesteps = overrides.pop( + "number_of_timesteps", self.number_of_timesteps + ) self.internal_arc.number_of_timesteps = self.number_of_timesteps self.internal_arc.decays.update(overrides.pop("decays", {})) super().apply_overrides(overrides)