X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/3fb451687293c993b10d860449eb405ddd44f377..9027ca63cac2b5ed24117ea6f5b6e760f1f2e11d:/plugin/black.vim

diff --git a/plugin/black.vim b/plugin/black.vim
index 0a26aa0..030c18f 100644
--- a/plugin/black.vim
+++ b/plugin/black.vim
@@ -37,10 +37,17 @@ if !exists("g:black_skip_string_normalization")
 endif
 
 python3 << endpython3
+import os
 import sys
 import vim
 
 def _get_python_binary(exec_prefix):
+  try:
+    default = vim.eval("g:pymode_python").strip()
+  except vim.error:
+    default = ""
+  if default and os.path.exists(default):
+    return default
   if sys.platform[:3] == "win":
     return exec_prefix / 'python.exe'
   return exec_prefix / 'bin' / 'python3'
@@ -87,8 +94,8 @@ def _initialize_black_env(upgrade=False):
     print('DONE! You are all set, thanks for waiting ✨ 🍰 ✨')
   if first_install:
     print('Pro-tip: to upgrade Black in the future, use the :BlackUpgrade command and restart Vim.\n')
-  if sys.path[0] != virtualenv_site_packages:
-    sys.path.insert(0, virtualenv_site_packages)
+  if virtualenv_site_packages not in sys.path:
+    sys.path.append(virtualenv_site_packages)
   return True
 
 if _initialize_black_env():
@@ -98,21 +105,49 @@ if _initialize_black_env():
 def Black():
   start = time.time()
   fast = bool(int(vim.eval("g:black_fast")))
-  line_length = int(vim.eval("g:black_linelength"))
-  mode = black.FileMode.AUTO_DETECT
-  if bool(int(vim.eval("g:black_skip_string_normalization"))):
-    mode |= black.FileMode.NO_STRING_NORMALIZATION
-  buffer_str = '\n'.join(vim.current.buffer) + '\n'
+  mode = black.FileMode(
+    line_length=int(vim.eval("g:black_linelength")),
+    string_normalization=not bool(int(vim.eval("g:black_skip_string_normalization"))),
+    is_pyi=vim.current.buffer.name.endswith('.pyi'),
+  )
+  (cursor_line, cursor_column) = vim.current.window.cursor
+  cb = vim.current.buffer[:]
+  cb_bc = cb[0:cursor_line]
+  # Format all code before the cursor.
+  # Detect unclosed blocks, close them with pass.
+  last_line = cb_bc[-1]
+  if last_line.rstrip().endswith(":"):
+      cb_bc[-1] = last_line + " pass"
+  # Determine old:new cursor location mapping
+  buffer_str_before = '\n'.join(cb_bc)+'\n'
+  try:
+    new_buffer_str_before = black.format_file_contents(buffer_str_before, fast=fast, mode=mode)
+    new_cb = new_buffer_str_before.split('\n')[:-1]
+    new_cursor_line = len(new_cb)
+    new_cursor = (new_cursor_line, cursor_column)
+  except black.NothingChanged:
+    new_cursor_line = cursor_line
+    new_cursor = (new_cursor_line, cursor_column)
+  except Exception as exc:
+    print(exc)
+  # Now we know where the cursor should be
+  # when we format the entire buffer. Do it:
+  buffer_str = '\n'.join(cb) + '\n'
   try:
-    new_buffer_str = black.format_file_contents(buffer_str, line_length=line_length, fast=fast, mode=mode)
+    new_buffer_str = black.format_file_contents(buffer_str, fast=fast, mode=mode)
   except black.NothingChanged:
     print(f'Already well formatted, good job. (took {time.time() - start:.4f}s)')
   except Exception as exc:
     print(exc)
   else:
-    cursor = vim.current.window.cursor
-    vim.current.buffer[:] = new_buffer_str.split('\n')[:-1]
-    vim.current.window.cursor = cursor
+    # Replace the buffer
+    new_cb = new_buffer_str.split('\n')[:-1]
+    vim.current.buffer[:] = new_cb
+    # Restore the cursor to its rightful place
+    try:
+      vim.current.window.cursor = new_cursor
+    except vim.error:
+      vim.current.window.cursor = (len(vim.current.buffer), 0)
     print(f'Reformatted in {time.time() - start:.4f}s.')
 
 def BlackUpgrade():