Added FrozenList subtraction

mcocdawc · Mar 1, 2017 · 73564ab · 73564ab
1 parent d0a281f
commit 73564ab
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/pandas/indexes/frozen.py b/pandas/indexes/frozen.py
@@ -28,8 +28,13 @@ def __add__(self, other):
         if isinstance(other, tuple):
             other = list(other)
         return self.__class__(super(FrozenList, self).__add__(other))
-
+    
     __iadd__ = __add__
+
+    def __sub__(self, other):
+        other = set(other)
+        temp = [x for x in self if x not in other]
+        return self.__class__(temp)
 
     # Python 2 compat
     def __getslice__(self, i, j):

diff --git a/pandas/tests/indexes/test_frozen.py b/pandas/tests/indexes/test_frozen.py
@@ -22,7 +22,17 @@ def test_add(self):
         result = (1, 2, 3) + self.container
         expected = FrozenList([1, 2, 3] + self.lst)
         self.check_result(result, expected)
-
+
+    def test_sub(self):
+        result = self.container - [2]
+        expected = FrozenList([1, 3, 4, 5])
+        self.check_result(result, expected) 
+
+    def test_sub_dupe(self):
+        result = FrozenList([1, 2, 3, 2]) - [2]
+        expected = FrozenList([1, 3])
+        self.check_result(result, expected) 
+
     def test_inplace(self):
         q = r = self.container
         q += [5]