Merge pull request #155 from octoenergy/finite-datetime-range-union

Add `union` method to `FiniteDatetimeRange`
octoenergy · Jun 17, 2024 · 80c65bc · 80c65bc
2 parents 163263f + 8e935b1
commit 80c65bc
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 0 deletions.
diff --git a/tests/test_ranges.py b/tests/test_ranges.py
@@ -887,6 +887,54 @@ def test_finite_range(self):
             assert 3 in subject
 
 
+class TestFiniteDatetimeRangeUnion:
+    def test_union_of_touching_ranges(self):
+        range = ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 1),
+            end=datetime.datetime(2000, 1, 2),
+        )
+        other = ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 2),
+            end=datetime.datetime(2000, 1, 3),
+        )
+
+        union = range | other
+
+        assert union == ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 1),
+            end=datetime.datetime(2000, 1, 3),
+        )
+
+    def test_union_of_disjoint_ranges(self):
+        range = ranges.FiniteDateRange(
+            start=datetime.datetime(2000, 1, 1),
+            end=datetime.datetime(2000, 1, 2),
+        )
+        other = ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2020, 1, 1),
+            end=datetime.datetime(2020, 1, 2),
+        )
+
+        assert range | other is None
+
+    def test_union_of_overlapping_ranges(self):
+        range = ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 1),
+            end=datetime.datetime(2000, 1, 3),
+        )
+        other = ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 2),
+            end=datetime.datetime(2000, 1, 4),
+        )
+
+        union = range | other
+
+        assert union == ranges.FiniteDatetimeRange(
+            start=datetime.datetime(2000, 1, 1),
+            end=datetime.datetime(2000, 1, 4),
+        )
+
+
 class TestAsFiniteDatetimePeriods:
     def test_converts(self):
         actually_finite_range = ranges.DatetimeRange(

diff --git a/xocto/ranges.py b/xocto/ranges.py
@@ -830,6 +830,23 @@ def intersection(
         assert base_intersection.boundaries == RangeBoundaries.INCLUSIVE_EXCLUSIVE
         return FiniteDatetimeRange(base_intersection.start, base_intersection.end)
 
+    def union(self, other: Range[datetime.datetime]) -> Optional["FiniteDatetimeRange"]:
+        """
+        Unions between two FiniteDatetimeRanges should produce a FiniteDatetimeRange.
+        """
+        try:
+            base_union = super().union(other)
+        except ValueError:
+            return None
+
+        if base_union is None:
+            return None
+
+        assert base_union.boundaries == RangeBoundaries.INCLUSIVE_EXCLUSIVE
+        assert base_union.start is not None
+        assert base_union.end is not None
+        return FiniteDatetimeRange(base_union.start, base_union.end)
+
     def __and__(
         self, other: Range[datetime.datetime]
     ) -> Optional["FiniteDatetimeRange"]: