added maxsplit argument to strutils.split

1 files changed, 25 insertions, 10 deletions

diff --git a/lib/pure/strutils.nim b/lib/pure/strutils.nim
index a78fed4b9..c9b23daf2 100644
--- a/lib/pure/strutils.nim
+++ b/lib/pure/strutils.nim
@@ -322,7 +322,7 @@ proc toOctal*(c: char): string {.noSideEffect, rtl, extern: "nsuToOctal".} =
     result[i] = chr(val mod 8 + ord('0'))
     val = val div 8
 
-iterator split*(s: string, seps: set[char] = Whitespace): string =
+iterator split*(s: string, seps: set[char] = Whitespace, maxsplit: int = -1): string =
   ## Splits the string `s` into substrings using a group of separators.
   ##
   ## Substrings are separated by a substring containing only `seps`. Note
@@ -367,15 +367,20 @@ iterator split*(s: string, seps: set[char] = Whitespace): string =
   ##   "08.398990"
   ##
   var last = 0
+  var splits = maxsplit
   assert(not ('\0' in seps))
   while last < len(s):
     while s[last] in seps: inc(last)
     var first = last
     while last < len(s) and s[last] notin seps: inc(last) # BUGFIX!
     if first <= last-1:
+      if splits == 0:
+        yield substr(s, first, len(s)-1)
+        break
       yield substr(s, first, last-1)
+      dec(splits)
 
-iterator split*(s: string, sep: char): string =
+iterator split*(s: string, sep: char, maxsplit: int = -1): string =
   ## Splits the string `s` into substrings using a single separator.
   ##
   ## Substrings are separated by the character `sep`.
@@ -402,26 +407,36 @@ iterator split*(s: string, sep: char): string =
   ##   ""
   ##
   var last = 0
+  var splits = maxsplit
   assert('\0' != sep)
   if len(s) > 0:
     # `<=` is correct here for the edge cases!
     while last <= len(s):
       var first = last
       while last < len(s) and s[last] != sep: inc(last)
+      if splits == 0:
+        yield substr(s, first, len(s)-1)
+        break
       yield substr(s, first, last-1)
+      dec(splits)
       inc(last)
 
-iterator split*(s: string, sep: string): string =
+iterator split*(s: string, sep: string, maxsplit: int = -1): string =
   ## Splits the string `s` into substrings using a string separator.
   ##
   ## Substrings are separated by the string `sep`.
   var last = 0
+  var splits = maxsplit
   if len(s) > 0:
     while last <= len(s):
       var first = last
       while last < len(s) and s.substr(last, last + <sep.len) != sep:
         inc(last)
+      if splits == 0:
+        yield substr(s, first, len(s)-1)
+        break
       yield substr(s, first, last-1)
+      dec(splits)
       inc(last, sep.len)
 
 iterator splitLines*(s: string): string =
@@ -491,25 +506,25 @@ proc countLines*(s: string): int {.noSideEffect,
     else: discard
     inc i
 
-proc split*(s: string, seps: set[char] = Whitespace): seq[string] {.
+proc split*(s: string, seps: set[char] = Whitespace, maxsplit: int = -1): seq[string] {.
   noSideEffect, rtl, extern: "nsuSplitCharSet".} =
   ## The same as the `split iterator <#split.i,string,set[char]>`_, but is a
   ## proc that returns a sequence of substrings.
-  accumulateResult(split(s, seps))
+  accumulateResult(split(s, seps, maxsplit))
 
-proc split*(s: string, sep: char): seq[string] {.noSideEffect,
+proc split*(s: string, sep: char, maxsplit: int = -1): seq[string] {.noSideEffect,
   rtl, extern: "nsuSplitChar".} =
   ## The same as the `split iterator <#split.i,string,char>`_, but is a proc
   ## that returns a sequence of substrings.
-  accumulateResult(split(s, sep))
+  accumulateResult(split(s, sep, maxsplit))
 
-proc split*(s: string, sep: string): seq[string] {.noSideEffect,
+proc split*(s: string, sep: string, maxsplit: int = -1): seq[string] {.noSideEffect,
   rtl, extern: "nsuSplitString".} =
   ## Splits the string `s` into substrings using a string separator.
   ##
   ## Substrings are separated by the string `sep`. This is a wrapper around the
   ## `split iterator <#split.i,string,string>`_.
-  accumulateResult(split(s, sep))
+  accumulateResult(split(s, sep, maxsplit))
 
 proc toHex*(x: BiggestInt, len: Positive): string {.noSideEffect,
   rtl, extern: "nsuToHex".} =
@@ -1660,7 +1675,7 @@ when isMainModule:
   doAssert isAlpha("Rasp")
   doAssert isAlpha("Args")
   doAssert(not isAlpha("$Tomato"))
-  
+
   doAssert isAlphaNumeric('3')
   doAssert isAlphaNumeric('R')
   doAssert(not isAlphaNumeric('!'))