X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/75a9447448d2e5d7acfa925bfead4a1460e1939d..d00eac9944acdafe2fb61a61c7e763c6502e89d5:/plugin/black.vim?ds=inline

diff --git a/plugin/black.vim b/plugin/black.vim
index 0b958b8..f0b55b2 100644
--- a/plugin/black.vim
+++ b/plugin/black.vim
@@ -2,7 +2,7 @@
 " Author: Łukasz Langa
 " Created: Mon Mar 26 23:27:53 2018 -0700
 " Requires: Vim Ver7.0+
-" Version:  1.0
+" Version:  1.1
 "
 " Documentation:
 "   This plugin formats Python files.
@@ -10,6 +10,8 @@
 " History:
 "  1.0:
 "    - initial version
+"  1.1:
+"    - restore cursor/window position after formatting
 
 if v:version < 700 || !has('python3')
     echo "This script requires vim7.0+ with Python 3.6 support."
@@ -100,6 +102,8 @@ def Black():
   mode = black.FileMode.AUTO_DETECT
   if bool(int(vim.eval("g:black_skip_string_normalization"))):
     mode |= black.FileMode.NO_STRING_NORMALIZATION
+  if vim.current.buffer.name.endswith('.pyi'):
+    mode |= black.FileMode.PYI
   buffer_str = '\n'.join(vim.current.buffer) + '\n'
   try:
     new_buffer_str = black.format_file_contents(buffer_str, line_length=line_length, fast=fast, mode=mode)
@@ -108,7 +112,12 @@ def Black():
   except Exception as exc:
     print(exc)
   else:
+    cursor = vim.current.window.cursor
     vim.current.buffer[:] = new_buffer_str.split('\n')[:-1]
+    try:
+      vim.current.window.cursor = cursor
+    except vim.error:
+      vim.current.window.cursor = (len(vim.current.buffer), 0)
     print(f'Reformatted in {time.time() - start:.4f}s.')
 
 def BlackUpgrade():
@@ -122,6 +131,3 @@ endpython3
 command! Black :py3 Black()
 command! BlackUpgrade :py3 BlackUpgrade()
 command! BlackVersion :py3 BlackVersion()
-
-nmap ,= :Black<CR>
-vmap ,= :Black<CR>