diff hglib/util.py @ 9:5882a698ad5c

util: rewrite eatlines (faster and simpler version)
author Idan Kamara <idankk86@gmail.com>
date Wed, 10 Aug 2011 01:48:38 +0300
parents 3ac38d500d68
children 19d2c55c3928
line wrap: on
line diff
--- a/hglib/util.py	Wed Aug 10 01:38:18 2011 +0300
+++ b/hglib/util.py	Wed Aug 10 01:48:38 2011 +0300
@@ -1,4 +1,4 @@
-import itertools
+import itertools, cStringIO
 
 def grouper(n, iterable):
     ''' list(grouper(2, range(4))) -> [(0, 1), (2, 3)] '''
@@ -6,11 +6,23 @@
     return itertools.izip(*args)
 
 def eatlines(s, n):
-    idx = 0
-    for i in xrange(n):
-        idx = s.find('\n', idx) + 1
+    """
+    >>> eatlines("1\\n2", 1)
+    '2'
+    >>> eatlines("1\\n2", 2)
+    ''
+    >>> eatlines("1\\n2", 3)
+    ''
+    >>> eatlines("1\\n2\\n3", 1)
+    '2\\n3'
+    """
+    cs = cStringIO.StringIO(s)
 
-    return s[idx:]
+    for line in cs:
+        n -= 1
+        if n == 0:
+            return cs.read()
+    return ''
 
 def cmdbuilder(name, *args, **kwargs):
     """