Commits

Sergey Astanin committed b412b3a

split() may now take function as the first argument

  • Participants
  • Parent commits ad5e5f6

Comments (0)

Files changed (1)

                 yield head
     return chopper()
 
-def _nextByDelim(delimiter, seq):
+def _nextByDelim(delimfunc, seq):
     "Next chunk from from the sequence seq, and sequence tail."
     iseq = iter(seq)
     chunk = []
     try:
         while True:
             x = next(iseq)
-            if x != delimiter:
+            if not delimfunc(x):
                 chunk.append(x)
             else:
                 break
 
 def split(delimiter, sequence, maxsplit=None):
     """
-    Break a sequence on elements equal to delimiter.
+    Break a sequence on particular elements.
     Return an iterator over chunks (delimiters excluded).
 
-    If maxsplit is given, at most maxsplit splits are done.
+    Arguments:
+
+    delimiter   if a function, it returns True on chunk separators;
+                otherwise, it is the value of chunk separator.
+    sequence    original sequence;
+    maxsplit    if given, at most maxsplit splits are done.
 
     >>> list(split(0, [1,2,3,0,4,5,0,0,6]))
     [[1, 2, 3], [4, 5], [], [6]]
     >>> list(map(list, split(0, [1,2,3,0,4,5,0,0,6], maxsplit=2)))
     [[1, 2, 3], [4, 5], [0, 6]]
 
+    >>> list(split(lambda x: x==5, range(10)))
+    [[0, 1, 2, 3, 4], [6, 7, 8, 9]]
+
     This function is lazy and produces new chunks only on demand:
 
     >>> if python_version_tuple()[0] > '2': xrange=range
     [0, 1, 2, 3, 4, 5, 6, 7, 8]
 
     """
+    if hasattr(delimiter, "__call__"):
+        delimfunc = delimiter
+    else:
+        delimfunc = lambda x: x == delimiter
     def splitter():
         tail = sequence
         splits = 0
                 yield tail
                 tail = None
             else:
-                chunk, tail = _nextByDelim(delimiter, tail)
+                chunk, tail = _nextByDelim(delimfunc, tail)
                 splits += 1
                 yield chunk
     return splitter()