add_value method (electricitymaps#5539)

* add_value method

* apply suggestion + format

* use super() in ProductionMix method override

* add tests

* format

electricitymap/contrib/lib/models/events.py

@@ -25,6 +25,24 @@ def set_value(self, mode: str, value: Optional[float]) -> None:
  """
  self.__setattr__(mode, value)
 
+ def add_value(
+ self,
+ mode: str,
+ value: Optional[float],
+ ) -> None:
+ """
+ Adds the provided value to the existing value of the provided mode.
+ This is useful if there are multiple production modes in the source
+ that maps to the same Electricity Maps production mode.
+ """
+ if value is None:
+ return
+ existing_value: Optional[float] = getattr(self, mode)
+ if existing_value is not None:
+ self.__setattr__(mode, existing_value + value)
+ else:
+ self.__setattr__(mode, value)
+
  @classmethod
  def merge(cls, mixes: List["Mix"]) -> "Mix":
  raise NotImplementedError()
@@ -98,7 +116,8 @@ def __setattr__(
  value: Optional[float],
  ) -> None:
  """
- Overriding the setattr method to check for negative values and set them to None.
+ Overriding the setattr method to check that the name is a valid production mode
+ and to check for negative values and set them to None.
  This method also keeps track of the modes that have been corrected.
  """
  if not name in PRODUCTION_MODES:
@@ -108,6 +127,18 @@ def __setattr__(
  value = None
  return super().__setattr__(name, value)
 
+ def _correct_negative_value(
+ self, mode: str, value: Optional[float], correct_negative_with_zero: bool
+ ) -> Union[float, None]:
+ """
+ Corrects a negative value by setting it to None or 0.
+ This method also keeps track of the modes that have been corrected.
+ """
+ if value is not None and value < 0:
+ self._corrected_negative_values.add(mode)
+ return 0 if correct_negative_with_zero else None
+ return value
+
  def set_value(
  self,
  mode: str,
@@ -119,11 +150,22 @@ def set_value(
  If correct_negative_with_zero is set to True, negative values will be set to 0 instead of None.
  This method keeps track of values that have been corrected.
  """
- if correct_negative_with_zero and value is not None and value < 0:
- value = 0
- self._corrected_negative_values.add(mode)
+ value = self._correct_negative_value(mode, value, correct_negative_with_zero)
  self.__setattr__(mode, value)
 
+ def add_value(
+ self,
+ mode: str,
+ value: Optional[float],
+ correct_negative_with_zero: bool = False,
+ ) -> None:
+ """Adds the provided value to the existing value of the provided mode. Negative values are set to None by default.
+ If correct_negative_with_zero is set to True, negative values will be set to 0 instead of None.
+ This method keeps track of modes that have been corrected.
+ """
+ value = self._correct_negative_value(mode, value, correct_negative_with_zero)
+ super().add_value(mode, value)
+
  @property
  def has_corrected_negative_values(self) -> bool:
  return len(self._corrected_negative_values) > 0

electricitymap/contrib/lib/tests/test_events.py

@@ -393,3 +393,66 @@ def test_storage_mix_has_all_storage_modes(self):
  mix = StorageMix()
  for mode in STORAGE_MODES:
  assert hasattr(mix, mode)
+
+
+class TestAddValue(unittest.TestCase):
+ def test_production(self):
+ mix = ProductionMix()
+ mix.add_value("wind", 10)
+ assert mix.wind == 10
+ mix.add_value("wind", 5)
+ assert mix.wind == 15
+ assert mix.corrected_negative_modes == set()
+
+ def test_production_with_negative_value(self):
+ mix = ProductionMix()
+ mix.add_value("wind", 10)
+ assert mix.wind == 10
+ mix.add_value("wind", -5)
+ assert mix.wind == 10
+ assert mix.corrected_negative_modes == set(["wind"])
+
+ def test_production_with_negative_value_expect_none(self):
+ mix = ProductionMix()
+ mix.add_value("wind", -10)
+ assert mix.wind == None
+ assert mix.corrected_negative_modes == set(["wind"])
+
+ def test_production_with_negative_value_and_correct_with_none(self):
+ mix = ProductionMix()
+ mix.add_value("wind", -10, correct_negative_with_zero=True)
+ assert mix.wind == 0
+ mix.add_value("wind", 15, correct_negative_with_zero=True)
+ assert mix.wind == 15
+ assert mix.corrected_negative_modes == set(["wind"])
+
+ def test_production_with_none(self):
+ mix = ProductionMix()
+ mix.add_value("wind", 10)
+ assert mix.wind == 10
+ mix.add_value("wind", None)
+ assert mix.wind == 10
+ assert mix.corrected_negative_modes == set()
+
+ def test_storage(self):
+ mix = StorageMix()
+ mix.add_value("hydro", 10)
+ assert mix.hydro == 10
+ mix.add_value("hydro", 5)
+ assert mix.hydro == 15
+
+ def test_storage_with_negative_value(self):
+ mix = StorageMix()
+ mix.add_value("hydro", 10)
+ assert mix.hydro == 10
+ mix.add_value("hydro", -5)
+ assert mix.hydro == 5
+
+ def test_storage_with_none(self):
+ mix = StorageMix()
+ mix.add_value("hydro", None)
+ assert mix.hydro == None
+ mix.add_value("hydro", -5)
+ assert mix.hydro == -5
+ mix.add_value("hydro", None)
+ assert mix.hydro == -5

pyproject.toml

@@ -105,7 +105,8 @@ load-plugins = "pylint_pydantic"
 [tool.pytest.ini_options]
 testpaths = [
  "tests",
- "parsers/test"
+ "parsers/test",
+ "electricitymap/contrib/lib/tests",
 ]
 
 [build-system]