Fix crop signature for python2

Toilal · Oct 18, 2015 · 1918fca · 1918fca
1 parent 7a9a7b8
commit 1918fca
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/rebulk/match.py b/rebulk/match.py
@@ -9,6 +9,7 @@
 
 #from .ordered_set import OrderedSet
 from .loose import ensure_list, filter_index
+from .utils import is_iterable
 
 
 class MatchesDict(dict):
@@ -489,14 +490,16 @@ def raw(self):
             return self.input_string[self.raw_start:self.raw_end]
         return None
 
-    def crop(self, *crops, predicate=None, index=None):
+    def crop(self, crops, predicate=None, index=None):
         """
         crop the match with given Match objects or spans tuples
         :param crops:
-        :type crops:
+        :type crops: list or object
         :return: a list of Match objects
         :rtype: list[Match]
         """
+        if not is_iterable(crops) or len(crops) == 2 and isinstance(crops[0], int):
+            crops = [crops]
         initial = copy.copy(self)
         ret = [initial]
         for crop in crops:

diff --git a/rebulk/test/test_match.py b/rebulk/test/test_match.py
@@ -161,7 +161,7 @@ def test_crop(self):
         match2 = Match(0, 2, input_string=input_string)
         match3 = Match(8, 15, input_string=input_string)
 
-        ret = match1.crop(match2, match3.span)
+        ret = match1.crop([match2, match3.span])
 
         assert len(ret) == 1
 
@@ -188,7 +188,7 @@ def test_crop(self):
         assert ret[0].span == (1, 4)
         assert ret[1].span == (6, 10)
 
-        ret = match1.crop((3, 5), (7, 9))
+        ret = match1.crop([(3, 5), (7, 9)])
         assert len(ret) == 3
 
         assert ret[0].span == (1, 3)