Add methods first_before and first_after.

intervaltree/intervaltree.py

@@ -762,6 +762,59 @@ def items(self):
         :rtype: set of Interval
         """
         return set(self.all_intervals)
+
+    def _first_after_or_after(self, value, key, optimum, compare):
+        node = self.top_node
+        first = None
+        if not node:
+            raise ValueError('Empty IntervalTree.')
+
+        def update_first(first, ivs):
+            try:
+                opt = optimum((iv for iv in ivs if compare(iv.begin, value)), key=key)
+            except ValueError:
+                pass
+            else:
+                if first is None or compare(key(first), key(opt)):
+                    first = opt
+            return first
+
+        if node.x_center < value:
+            while node is not None and not compare(node.x_center, value):
+                first = update_first(first, node.s_center)
+                node = node.right_node
+        else:
+            while node is not None and compare(node.x_center, value):
+                first = update_first(first, node.s_center)
+                node = node.left_node
+        if node is None:
+            if first:
+                return first
+            else:
+                raise ValueError('No interval before/after {0}.'.format(value))
+        else:
+            candidates = [iv for iv in node.all_children() if compare(key(iv), value)]
+            if first:
+                candidates.append(first)
+            if not candidates:
+                raise ValueError('No interval before/after {0}.'.format(value))
+            return optimum(candidates, key=key)
+
+    def first_after(self, value):
+        """Return an interval whose beginning is ≥ value and is minimal."""
+        return self._first_after_or_after(value,
+                key=lambda iv: iv.begin,
+                optimum=min,
+                compare=lambda x, y: x >= y)
+
+    def first_before(self, value):
+        """Return an interval whose end is ≤ value and is maximal."""
+        return self._first_after_or_after(value,
+                key=lambda iv: iv.end,
+                optimum=max,
+                compare=lambda x, y: x <= y)
+
+
 
     def is_empty(self):
         """

test/intervaltree_methods/first_before_after_test.py

@@ -0,0 +1,45 @@
+from __future__ import absolute_import
+from intervaltree import Interval, IntervalTree
+import pytest
+from test.intervaltrees import trees, sdata
+
+def test_first_after():
+    t = trees['ivs1']()
+    assert t.first_after(7) in [
+            Interval(8, 10, '[8,10)'),
+            Interval(8, 15, '[8,15)'),
+            ]
+    assert t.first_after(8) in [
+            Interval(8, 10, '[8,10)'),
+            Interval(8, 15, '[8,15)'),
+            ]
+    assert t.first_after(13) == Interval(14, 15, '[14,15)')
+    assert t.first_after(4) == Interval(4, 7, '[4,7)')
+    assert t.first_after(5) == Interval(5, 9, '[5,9)')
+    assert t.first_after(-100) == Interval(1, 2, '[1,2)')
+    try:
+        iv = t.first_after(15)
+    except Exception as e:
+        assert isinstance(e, ValueError)
+    else:
+        assert False, iv
+
+def test_first_before():
+    t = trees['ivs1']()
+    assert t.first_before(5) == Interval(1, 2, '[1,2)')
+    assert t.first_before(7) == Interval(4, 7, '[4,7)')
+    assert t.first_before(10) in [
+            Interval(6, 10, '[6,10)'),
+            Interval(8, 10, '[8,10)'),
+            ]
+    assert t.first_before(15) == Interval(14, 15, '[14,15)')
+    assert t.first_before(100) == Interval(14, 15, '[14,15)')
+    try:
+        iv = t.first_before(1)
+    except Exception as e:
+        assert isinstance(e, ValueError)
+    else:
+        assert False, iv
+
+if __name__ == "__main__":
+    pytest.main([__file__, '-v'])