Add contains function to Bounds and SystemBounds

Signed-off-by: Jack <[email protected]>

Author: jack-herrmann

src/frequenz/sdk/timeseries/_base_types.py

@@ -8,7 +8,7 @@
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Generic, Self, TypeVar, overload
+from typing import Any, Generic, Protocol, Self, TypeVar, overload, runtime_checkable
 
 from ._quantities import Power, QuantityT
 
@@ -134,6 +134,39 @@ def map(
         )
 
 
+@runtime_checkable
+class Comparable(Protocol):
+    """
+    A protocol that requires the implementation of comparison methods.
+
+    This protocol is used to ensure that types can be compared using
+    the less than or equal to (`<=`) and greater than or equal to (`>=`)
+    operators.
+    """
+
+    def __le__(self, other: Any, /) -> bool:
+        """
+        Return True if self is less than or equal to other, otherwise False.
+
+        Args:
+            other: The value to compare with.
+
+        Returns:
+            bool: True if self is less than or equal to other, otherwise False.
+        """
+
+    def __ge__(self, other: Any, /) -> bool:
+        """
+        Return True if self is greater than or equal to other, otherwise False.
+
+        Args:
+            other: The value to compare with.
+
+        Returns:
+            bool: True if self is greater than or equal to other, otherwise False.
+        """
+
+
 _T = TypeVar("_T")
 
 
@@ -147,6 +180,29 @@ class Bounds(Generic[_T]):
     upper: _T
     """Upper bound."""
 
+    def __contains__(self, item: _T) -> bool:
+        """
+        Check if the value is within the range of the container.
+
+        Args:
+            item: The value to check.
+
+        Returns:
+            bool: True if value is within the range, otherwise False.
+        """
+        if item is None:
+            return False
+
+        assert isinstance(item, Comparable)
+
+        if self.lower is not None and self.upper is not None:
+            return self.lower <= item <= self.upper
+        if self.lower is not None:
+            return self.lower <= item
+        if self.upper is not None:
+            return item <= self.upper
+        return False
+
 
 @dataclass(frozen=True, kw_only=True)
 class SystemBounds:
@@ -171,3 +227,21 @@ class SystemBounds:
     This is the range within which power requests are NOT allowed by the pool.
     If present, they will be a subset of the inclusion bounds.
     """
+
+    def __contains__(self, item: Power) -> bool:
+        """
+        Check if the value is within the range of the container.
+
+        Args:
+            item: The value to check.
+
+        Returns:
+            bool: True if value is within the range, otherwise False.
+        """
+        if not self.inclusion_bounds:
+            return False
+        if item not in self.inclusion_bounds:
+            return False
+        if self.exclusion_bounds and item in self.exclusion_bounds:
+            return False
+        return True

tests/timeseries/test_base_types.py

@@ -0,0 +1,79 @@
+# License: MIT
+# Copyright © 2023 Frequenz Energy-as-a-Service GmbH
+
+"""Tests for battery pool."""
+
+
+from datetime import datetime
+
+from frequenz.sdk.timeseries._base_types import Bounds, SystemBounds
+from frequenz.sdk.timeseries._quantities import Power
+
+
+def test_bounds_contains() -> None:
+    """
+    Test the `__contains__` method of the `Bounds` class.
+
+    This method checks if a value is within the defined bounds.
+    """
+    bounds = Bounds(lower=Power.from_watts(10), upper=Power.from_watts(100))
+
+    assert Power.from_watts(50) in bounds  # within bounds
+    assert Power.from_watts(10) in bounds  # at lower bound
+    assert Power.from_watts(100) in bounds  # at upper bound
+    assert Power.from_watts(9) not in bounds  # below lower bound
+    assert Power.from_watts(101) not in bounds  # above upper bound
+
+    bounds_no_lower = Bounds(lower=None, upper=Power.from_watts(100))
+    assert Power.from_watts(50) in bounds_no_lower  # within upper bound
+    assert Power.from_watts(100) in bounds_no_lower  # at upper bound
+    assert Power.from_watts(101) not in bounds_no_lower  # above upper bound
+
+    bounds_no_upper = Bounds(lower=Power.from_watts(10), upper=None)
+    assert Power.from_watts(50) in bounds_no_upper  # within lower bound
+    assert Power.from_watts(10) in bounds_no_upper  # at lower bound
+    assert Power.from_watts(9) not in bounds_no_upper  # below lower bound
+
+    # test succeeds, mypy 'Unsupported operand types for in ("Power" and "Bounds[None]")'
+    # bounds_none = Bounds(lower=None, upper=None)
+    # assert Power.from_watts(15) not in bounds_none
+
+
+def test_system_bounds_contains() -> None:
+    """
+    Test the `__contains__` method of the `SystemBounds` class.
+
+    Checks if value is within inclusion bounds and not exclusion bounds.
+    """
+    inclusion_bounds = Bounds(lower=Power.from_watts(10), upper=Power.from_watts(100))
+    exclusion_bounds = Bounds(lower=Power.from_watts(40), upper=Power.from_watts(50))
+    system_bounds = SystemBounds(
+        timestamp=datetime.now(),
+        inclusion_bounds=inclusion_bounds,
+        exclusion_bounds=exclusion_bounds,
+    )
+
+    assert Power.from_watts(30) in system_bounds  # within inclusion, not in exclusion
+    assert (
+        Power.from_watts(45) not in system_bounds
+    )  # within inclusion, also in exclusion
+    assert Power.from_watts(110) not in system_bounds  # outside inclusion bounds
+
+    system_bounds_no_exclusion = SystemBounds(
+        timestamp=datetime.now(),
+        inclusion_bounds=inclusion_bounds,
+        exclusion_bounds=None,
+    )
+    assert Power.from_watts(30) in system_bounds_no_exclusion  # within inclusion bounds
+    assert (
+        Power.from_watts(110) not in system_bounds_no_exclusion
+    )  # outside inclusion bounds
+
+    system_bounds_no_inclusion = SystemBounds(
+        timestamp=datetime.now(),
+        inclusion_bounds=None,
+        exclusion_bounds=exclusion_bounds,
+    )
+    assert (
+        Power.from_watts(30) not in system_bounds_no_inclusion
+    )  # not in inclusion bounds