Handle iterable operands consistently in set operations

intervaltree/intervaltree.py

@@ -29,9 +29,9 @@
 from copy import copy
 
 try:
-    from collections.abc import MutableSet  # Python 3?
+    from collections.abc import MutableSet, Set as AbstractSet  # Python 3?
 except ImportError:
-    from collections import MutableSet
+    from collections import MutableSet, Set as AbstractSet
 
 try:
     xrange  # Python 2?
@@ -402,6 +402,8 @@ def difference(self, other):
         Returns a new tree, comprising all intervals in self but not
         in other.
         """
+        if not isinstance(other, AbstractSet):
+            other = set(other)
         ivs = set()
         for iv in self:
             if iv not in other:
@@ -412,6 +414,7 @@ def difference_update(self, other):
         """
         Removes all intervals in other from self.
         """
+        other = set(other)
         for iv in other:
             self.discard(iv)
 
@@ -427,6 +430,8 @@ def intersection(self, other):
         Returns a new tree of all intervals common to both self and
         other.
         """
+        if not isinstance(other, AbstractSet):
+            other = set(other)
         ivs = set()
         shorter, longer = sorted([self, other], key=len)
         for iv in shorter:
@@ -438,6 +443,8 @@ def intersection_update(self, other):
         """
         Removes intervals from self unless they also exist in other.
         """
+        if not isinstance(other, AbstractSet):
+            other = set(other)
         ivs = list(self)
         for iv in ivs:
             if iv not in other:

test/intervaltree_methods/setlike_test.py

@@ -21,6 +21,10 @@
 from __future__ import absolute_import
 from intervaltree import Interval, IntervalTree
 import pytest
+try:
+    from collections.abc import Set as AbstractSet
+except ImportError:
+    from collections import Set as AbstractSet
 from test import data
 try:
     import cPickle as pickle
@@ -461,3 +465,109 @@ def test_symmetric_difference():
 
 if __name__ == "__main__":
     pytest.main([__file__, '-v'])
+
+
+@pytest.mark.parametrize("method", [
+    "difference", "intersection", "intersection_update",
+])
+@pytest.mark.parametrize("operand_type", [list, tuple, set, frozenset, IntervalTree, iter])
+def test_set_operations_accept_iterable_operands(method, operand_type):
+    tree = IntervalTree.from_tuples([(0, 1), (2, 3), (4, 5)])
+    original = set(tree)
+    # Reversed iteration forces membership checks to reuse the operand.
+    operand = operand_type(list(reversed(list(tree))))
+    expected = (original.difference(original) if method == "difference"
+                else original.intersection(original))
+
+    result = getattr(tree, method)(operand)
+
+    if method == "intersection_update":
+        assert result is None
+        result = tree
+    else:
+        assert set(tree) == original
+    assert isinstance(result, IntervalTree)
+    assert set(result) == expected
+    result.verify()
+
+
+@pytest.mark.parametrize("method", [
+    "difference", "intersection", "intersection_update",
+])
+def test_set_operations_accept_generator_operands(method):
+    tree = IntervalTree.from_tuples([(0, 1), (2, 3), (4, 5)])
+    original = set(tree)
+    other = set([Interval(2, 3), Interval(6, 7)])
+    operand = (iv for iv in other)
+    expected = (original.difference(other) if method == "difference"
+                else original.intersection(other))
+
+    result = getattr(tree, method)(operand)
+
+    if method == "intersection_update":
+        assert result is None
+        result = tree
+    else:
+        assert set(tree) == original
+    assert set(result) == expected
+    result.verify()
+
+
+def test_difference_update_with_self():
+    tree = IntervalTree.from_tuples([(0, 1), (2, 3), (4, 5)])
+
+    assert tree.difference_update(tree) is None
+
+    assert not tree
+    tree.verify()
+
+
+class CountingAbstractSet(AbstractSet):
+    def __init__(self, items):
+        self.items = set(items)
+        self.contains_calls = 0
+        self.iter_calls = 0
+
+    def __len__(self):
+        return len(self.items)
+
+    def __contains__(self, item):
+        self.contains_calls += 1
+        return item in self.items
+
+    def __iter__(self):
+        self.iter_calls += 1
+        return iter(self.items)
+
+
+class CountingIntervalTree(IntervalTree):
+    def __init__(self, items):
+        self.iter_calls = 0
+        super(CountingIntervalTree, self).__init__(items)
+
+    def __iter__(self):
+        self.iter_calls += 1
+        return super(CountingIntervalTree, self).__iter__()
+
+
+@pytest.mark.parametrize("method", [
+    "difference", "intersection", "intersection_update",
+])
+@pytest.mark.parametrize("operand_type", [
+    CountingAbstractSet, CountingIntervalTree,
+])
+def test_set_operations_preserve_set_operands(method, operand_type):
+    interval = Interval(0, 1)
+    tree = IntervalTree([interval])
+    other = operand_type([interval, Interval(2, 3)])
+
+    result = getattr(tree, method)(other)
+
+    if method == "intersection_update":
+        result = tree
+    expected = set() if method == "difference" else set([interval])
+    assert set(result) == expected
+    assert other.iter_calls == 0
+    if isinstance(other, CountingAbstractSet):
+        assert other.contains_calls == 1
+    result.verify()