Merge pull request #18 from equalsraf/tb-pydo

Fix pydo

Author: tarruda

neovim/plugins/script_host.py

@@ -3,11 +3,14 @@
 import sys
 
 from ..api.common import SessionHook
+from ..compat import IS_PYTHON3
 
 
 logger = logging.getLogger(__name__)
 debug, warn = (logger.debug, logger.warn,)
 
+if IS_PYTHON3:
+    basestring = str
 
 class ScriptHost(object):
     """
@@ -55,13 +58,31 @@ def python_do_range(self, start, stop, code):
             sstart = start
             sstop = min(start + 5000, stop)
             lines = nvim.current.buffer.get_line_slice(sstart, sstop, True, True)
+
+            exception = None
+            newlines = []
+            linenr = sstart + 1
             for i, line in enumerate(lines):
-                linenr = i + sstart + 1
-                result = str(function(line, linenr))
-                if result:
-                    lines[i] = result
+                result = function(line, linenr)
+                if result == None:
+                    # Update earlier lines, and skip to the next
+                    if newlines:
+                        nvim.current.buffer.set_line_slice(sstart, sstart + len(newlines) -1, True, True, newlines)
+                    sstart += len(newlines) + 1
+                    newlines = []
+                    pass
+                elif isinstance(result,basestring):
+                    newlines.append(result)
+                else:
+                    exception = TypeError('pydo should return a string or None, found %s instead' % result.__class__.__name__)
+                    break
+                linenr += 1
+
             start = sstop + 1
-            nvim.current.buffer.set_line_slice(sstart, sstop, True, True, lines)
+            if newlines:
+                nvim.current.buffer.set_line_slice(sstart, sstart + len(newlines)-1, True, True, newlines)
+            if exception:
+                raise exception
         # delete the function
         del self.module.__dict__[fname]
 

test/test_script_host.py

@@ -56,11 +56,17 @@ def test_pyfile():
 
 @with_setup(setup=host_setup, teardown=host_teardown)
 def test_pydo():
+
+    # :pydo 42 returns None for all lines,
+    # the buffer should not be changed
+    vim.command('normal :pydo 42\n')
+    eq(vim.current.buffer.options['mod'], False)
+
     # insert some text
     vim.command('normal iabc\ndef\nghi')
     eq(vim.current.buffer[:], ['abc', 'def', 'ghi'])
     # go to top and select and replace the first two lines
-    vim.command('normal ggvj:pydo return linenr\n')
+    vim.command('normal ggvj:pydo return str(linenr)\n')
     eq(vim.current.buffer[:], ['1', '2', 'ghi'])