Bugfixes for floor and ceiling

test_treeset.py

@@ -50,6 +50,10 @@ def test_treeset(self):
 
         ts.clear()
         self.assertEqual(ts._treeset, [])
+
+        ts = TreeSet([int(x) for x in '934853458236'])
+        self.assertEqual(ts.ceiling(100), None)
+        self.assertEqual(ts.floor(0), None)
 
 if __name__ == '__main__':
     unittest.main()

treeset.py

@@ -32,14 +32,19 @@ def ceiling(self, e):
         index = bisect.bisect_right(self._treeset, e)
         if self[index - 1] == e:
             return e
-        return self._treeset[bisect.bisect_right(self._treeset, e)]
+        try:
+            return self._treeset[bisect.bisect_right(self._treeset, e)]
+        except IndexError:
+            return None
 
     def floor(self, e):
         index = bisect.bisect_left(self._treeset, e)
         if self[index] == e:
             return e
-        else:
-            return self._treeset[bisect.bisect_left(self._treeset, e) - 1]
+        check = self._treeset[bisect.bisect_left(self._treeset, e) - 1]
+        if check <= e:
+            return check
+        return None
 
     def __getitem__(self, num):
         return self._treeset[num]
@@ -48,31 +53,23 @@ def __len__(self):
         return len(self._treeset)
 
     def clear(self):
-        """
-        Delete all elements in TreeSet.
-        """
+        """Delete all elements in TreeSet."""
         self._treeset = []
 
     def clone(self):
-        """
-        Return shallow copy of self.
-        """
+        """Return shallow copy of self."""
         return TreeSet(self._treeset)
 
     def remove(self, element):
-        """
-        Remove element if element in TreeSet.
-        """
+        """Remove element if element in TreeSet."""
         try:
             self._treeset.remove(element)
         except ValueError:
             return False
         return True
 
     def __iter__(self):
-        """
-        Do ascending iteration for TreeSet
-        """
+        """Do ascending iteration for TreeSet"""
         for element in self._treeset:
             yield element
 
@@ -87,14 +84,13 @@ def __eq__(self, target):
             return self._treeset == target.treeset
         elif isinstance(target, list):
             return self._treeset == target
+        return None
 
     def __contains__(self, e):
-        """
-        Fast attribution judgment by bisect
-        """
+        """Fast attribution judgment by bisect"""
         try:
             return e == self._treeset[bisect.bisect_left(self._treeset, e)]
-        except:
+        except Exception:
             return False
 
 if __name__ == '__main__':