implements union

indexOutOfBound5 · Apr 19, 2017 · 523e8ac · 523e8ac
1 parent 23e58ec
commit 523e8ac
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 7 deletions.
diff --git a/python/common/org/python/types/FrozenSet.java b/python/common/org/python/types/FrozenSet.java
@@ -324,16 +324,16 @@ public org.python.Object __xor__(org.python.Object other) {
         throw new org.python.exceptions.TypeError("unsupported operand type(s) for ^: '" + this.typeName() + "' and '" + other.typeName() + "'");
     }
 
-    private org.python.types.FrozenSet iterToFrozenSet(org.python.Object iterable) {
+    private java.util.Set iterToSet(org.python.Object iterable) {
         org.python.Iterable iterator = iterable.__iter__();
-        java.util.Set<org.python.Object> temp = new java.util.HashSet<org.python.Object>();
+        java.util.Set<org.python.Object> set = new java.util.HashSet<org.python.Object>();
         try {
             while (true) {
-                temp.add(iterator.__next__());
+                set.add(iterator.__next__());
             }
         } catch (org.python.exceptions.StopIteration si) {
         }
-        return new org.python.types.FrozenSet(temp);
+        return set;
     }
 
     @org.python.Method(
@@ -344,7 +344,7 @@ private org.python.types.FrozenSet iterToFrozenSet(org.python.Object iterable) {
     public org.python.Object isdisjoint(org.python.Object other) {
         try {
             if (!(other instanceof org.python.types.Set || other instanceof org.python.types.FrozenSet)) {
-                other = iterToFrozenSet(other);
+                other = new org.python.types.FrozenSet(iterToSet(other));
             }
             org.python.types.FrozenSet temp = (org.python.types.FrozenSet) this.__and__(other);
             if (temp.__len__().value > 0) {
@@ -364,7 +364,7 @@ public org.python.Object isdisjoint(org.python.Object other) {
     public org.python.Object issubset(org.python.Object other) {
         try {
             if (!(other instanceof org.python.types.Set || other instanceof org.python.types.FrozenSet)) {
-                other = iterToFrozenSet(other);
+                other = new org.python.types.FrozenSet(iterToSet(other));
             }
             return this.__le__(other);
         } catch (org.python.exceptions.AttributeError e) {
@@ -379,11 +379,35 @@ public org.python.Object issubset(org.python.Object other) {
     public org.python.Object issuperset(org.python.Object other) {
         try {
             if (!(other instanceof org.python.types.Set || other instanceof org.python.types.FrozenSet)) {
-                other = iterToFrozenSet(other);
+                other = new org.python.types.FrozenSet(iterToSet(other));
             }
             return this.__ge__(other);
         } catch (org.python.exceptions.AttributeError e) {
             throw new org.python.exceptions.TypeError("'" + other.typeName() + "' object is not iterable");
         }
     }
+
+    @org.python.Method(
+            __doc__ = "Return the union of sets as a new set.\n\n(i.e. all elements that are in either set.)",
+            varargs = "others"
+    )
+    public org.python.Object union(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.addAll(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
@@ -216,6 +216,27 @@ def test_issuperset(self):
                 print(err)
             """)
 
+    def test_union(self):
+        self.assertCodeExecution("""
+            x = frozenset({1, 2, 3})
+            y = frozenset({3, 4, 5})
+            z = [5, 6, 7]
+            w = 1
+
+            print(x.union(y))
+
+            # multiple args test
+            print(x.union(y, z))
+
+            # iterable test
+            print(x.union(z))
+
+            # not-iterable test
+            try:
+                print(x.union(w))
+            except TypeError as err:
+                print(err)
+            """)
 
 class UnaryFrozensetOperationTests(UnaryOperationTestCase, TranspileTestCase):
     data_type = 'frozenset'