implements difference

indexOutOfBound5 · Apr 19, 2017 · a4dbb48 · a4dbb48
1 parent 3fee4d6
commit a4dbb48
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 0 deletions.
diff --git a/python/common/org/python/types/FrozenSet.java b/python/common/org/python/types/FrozenSet.java
@@ -434,4 +434,28 @@ public org.python.Object intersection(org.python.types.Tuple others) {
         }
         return new org.python.types.FrozenSet(set);
     }
+
+    @org.python.Method(
+            __doc__ = "Return the difference of two or more sets as a new set.\n\n(i.e. all elements that are in this set but not the others.)",
+            varargs = "others"
+    )
+    public org.python.Object difference(org.python.types.Tuple others) {
+        java.util.Set set = new java.util.HashSet<org.python.Object>(this.value);
+        for (org.python.Object other: others.value) {
+            try {
+                java.util.Set otherSet = null;
+                if (other instanceof org.python.types.Set) {
+                    otherSet = ((org.python.types.Set) other).value;
+                } else if (other instanceof org.python.types.FrozenSet) {
+                    otherSet = ((org.python.types.FrozenSet) other).value;
+                } else {
+                    otherSet = iterToSet(other);
+                }
+                set.removeAll(otherSet);
+            } catch (org.python.exceptions.AttributeError e) {
+                throw new org.python.exceptions.TypeError("'" + other.typeName() + "' object is not iterable");
+            }
+        }
+        return new org.python.types.FrozenSet(set);
+    }
 }
diff --git a/tests/datatypes/test_frozenset.py b/tests/datatypes/test_frozenset.py
@@ -268,6 +268,33 @@ def test_intersection(self):
                 print(err)
             """)
 
+    def test_difference(self):
+        self.assertCodeExecution("""
+            x = frozenset({1, 2, 3})
+            y = frozenset({2, 3, 4})
+            z = [1, 6]
+            w = 1
+            t = frozenset()
+
+            print(x.difference(y))
+
+            # empty set test
+            print(x.difference(t))
+
+            # multiple args test
+            print(x.difference(y, z))
+
+            # iterable test
+            print(x.difference(z))
+
+            # not-iterable test
+            try:
+                print(x.difference(w))
+            except TypeError as err:
+                print(err)
+            """)
+
+
 class UnaryFrozensetOperationTests(UnaryOperationTestCase, TranspileTestCase):
     data_type = 'frozenset'