implements intersection

trillium · Apr 19, 2017 · 3fee4d6 · 3fee4d6
1 parent 523e8ac
commit 3fee4d6
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/python/common/org/python/types/FrozenSet.java b/python/common/org/python/types/FrozenSet.java
@@ -410,4 +410,28 @@ public org.python.Object union(org.python.types.Tuple others) {
         }
         return new org.python.types.FrozenSet(set);
     }
+
+    @org.python.Method(
+            __doc__ = "Return the intersection of two sets as a new set.\n\n(i.e. all elements that are in both sets.)",
+            varargs = "others"
+    )
+    public org.python.Object intersection(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.retainAll(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
@@ -222,9 +222,13 @@ def test_union(self):
             y = frozenset({3, 4, 5})
             z = [5, 6, 7]
             w = 1
+            t = frozenset()
 
             print(x.union(y))
 
+            # empty set test
+            print(x.union(t))
+
             # multiple args test
             print(x.union(y, z))
 
@@ -238,6 +242,32 @@ def test_union(self):
                 print(err)
             """)
 
+    def test_intersection(self):
+        self.assertCodeExecution("""
+            x = frozenset({1, 2, 3})
+            y = frozenset({2, 3, 4})
+            z = [3, 6, 7]
+            w = 1
+            t = frozenset()
+
+            print(x.intersection(y))
+
+            # empty set test
+            print(x.intersection(t))
+
+            # multiple args test
+            print(x.intersection(y, z))
+
+            # iterable test
+            print(x.intersection(z))
+
+            # not-iterable test
+            try:
+                print(x.intersection(w))
+            except TypeError as err:
+                print(err)
+            """)
+
 class UnaryFrozensetOperationTests(UnaryOperationTestCase, TranspileTestCase):
     data_type = 'frozenset'