plugin/black.vim

   1 " black.vim
   2 " Author: Łukasz Langa
   3 " Created: Mon Mar 26 23:27:53 2018 -0700
   4 " Requires: Vim Ver7.0+
   5 " Version:  1.1
   6 "
   7 " Documentation:
   8 "   This plugin formats Python files.
   9 "
  10 " History:
  11 "  1.0:
  12 "    - initial version
  13 "  1.1:
  14 "    - restore cursor/window position after formatting
  15
  16 if v:version < 700 || !has('python3')
  17     echo "This script requires vim7.0+ with Python 3.6 support."
  18     finish
  19 endif
  20
  21 if exists("g:load_black")
  22    finish
  23 endif
  24
  25 let g:load_black = "py1.0"
  26 if !exists("g:black_virtualenv")
  27   if has("nvim")
  28     let g:black_virtualenv = "~/.local/share/nvim/black"
  29   else
  30     let g:black_virtualenv = "~/.vim/black"
  31   endif
  32 endif
  33 if !exists("g:black_fast")
  34   let g:black_fast = 0
  35 endif
  36 if !exists("g:black_linelength")
  37   let g:black_linelength = 88
  38 endif
  39 if !exists("g:black_skip_string_normalization")
  40   let g:black_skip_string_normalization = 0
  41 endif
  42
  43 python3 << endpython3
  44 import os
  45 import sys
  46 import vim
  47
  48 def _get_python_binary(exec_prefix):
  49   try:
  50     default = vim.eval("g:pymode_python").strip()
  51   except vim.error:
  52     default = ""
  53   if default and os.path.exists(default):
  54     return default
  55   if sys.platform[:3] == "win":
  56     return exec_prefix / 'python.exe'
  57   return exec_prefix / 'bin' / 'python3'
  58
  59 def _get_pip(venv_path):
  60   if sys.platform[:3] == "win":
  61     return venv_path / 'Scripts' / 'pip.exe'
  62   return venv_path / 'bin' / 'pip'
  63
  64 def _get_virtualenv_site_packages(venv_path, pyver):
  65   if sys.platform[:3] == "win":
  66     return venv_path / 'Lib' / 'site-packages'
  67   return venv_path / 'lib' / f'python{pyver[0]}.{pyver[1]}' / 'site-packages'
  68
  69 def _initialize_black_env(upgrade=False):
  70   pyver = sys.version_info[:2]
  71   if pyver < (3, 6):
  72     print("Sorry, Black requires Python 3.6+ to run.")
  73     return False
  74
  75   from pathlib import Path
  76   import subprocess
  77   import venv
  78   virtualenv_path = Path(vim.eval("g:black_virtualenv")).expanduser()
  79   virtualenv_site_packages = str(_get_virtualenv_site_packages(virtualenv_path, pyver))
  80   first_install = False
  81   if not virtualenv_path.is_dir():
  82     print('Please wait, one time setup for Black.')
  83     _executable = sys.executable
  84     try:
  85       sys.executable = str(_get_python_binary(Path(sys.exec_prefix)))
  86       print(f'Creating a virtualenv in {virtualenv_path}...')
  87       print('(this path can be customized in .vimrc by setting g:black_virtualenv)')
  88       venv.create(virtualenv_path, with_pip=True)
  89     finally:
  90       sys.executable = _executable
  91     first_install = True
  92   if first_install:
  93     print('Installing Black with pip...')
  94   if upgrade:
  95     print('Upgrading Black with pip...')
  96   if first_install or upgrade:
  97     subprocess.run([str(_get_pip(virtualenv_path)), 'install', '-U', 'black'], stdout=subprocess.PIPE)
  98     print('DONE! You are all set, thanks for waiting ✨ 🍰 ✨')
  99   if first_install:
 100     print('Pro-tip: to upgrade Black in the future, use the :BlackUpgrade command and restart Vim.\n')
 101   if virtualenv_site_packages not in sys.path:
 102     sys.path.append(virtualenv_site_packages)
 103   return True
 104
 105 if _initialize_black_env():
 106   import black
 107   import time
 108
 109 def Black():
 110   start = time.time()
 111   fast = bool(int(vim.eval("g:black_fast")))
 112   mode = black.FileMode(
 113     line_length=int(vim.eval("g:black_linelength")),
 114     string_normalization=not bool(int(vim.eval("g:black_skip_string_normalization"))),
 115     is_pyi=vim.current.buffer.name.endswith('.pyi'),
 116   )
 117   (cursor_line, cursor_column) = vim.current.window.cursor
 118   cb = vim.current.buffer[:]
 119   cb_bc = cb[0:cursor_line]
 120   # Format all code before the cursor.
 121   # Detect unclosed blocks, close them with pass.
 122   last_line = cb_bc[-1]
 123   if last_line.rstrip().endswith(":"):
 124       cb_bc[-1] = last_line + " pass"
 125   # Determine old:new cursor location mapping
 126   buffer_str_before = '\n'.join(cb_bc)+'\n'
 127   try:
 128     new_buffer_str_before = black.format_file_contents(buffer_str_before, fast=fast, mode=mode)
 129     new_cb = new_buffer_str_before.split('\n')[:-1]
 130     new_cursor_line = len(new_cb)
 131     new_cursor = (new_cursor_line, cursor_column)
 132   except black.NothingChanged:
 133     new_cursor_line = cursor_line
 134     new_cursor = (new_cursor_line, cursor_column)
 135   except Exception as exc:
 136     print(exc)
 137   # Now we know where the cursor should be
 138   # when we format the entire buffer. Do it:
 139   buffer_str = '\n'.join(cb) + '\n'
 140   try:
 141     new_buffer_str = black.format_file_contents(buffer_str, fast=fast, mode=mode)
 142   except black.NothingChanged:
 143     print(f'Already well formatted, good job. (took {time.time() - start:.4f}s)')
 144   except Exception as exc:
 145     print(exc)
 146   else:
 147     # Replace the buffer
 148     new_cb = new_buffer_str.split('\n')[:-1]
 149     vim.current.buffer[:] = new_cb
 150     # Restore the cursor to its rightful place
 151     try:
 152       vim.current.window.cursor = new_cursor
 153     except vim.error:
 154       vim.current.window.cursor = (len(vim.current.buffer), 0)
 155     print(f'Reformatted in {time.time() - start:.4f}s.')
 156
 157 def BlackUpgrade():
 158   _initialize_black_env(upgrade=True)
 159
 160 def BlackVersion():
 161   print(f'Black, version {black.__version__} on Python {sys.version}.')
 162
 163 endpython3
 164
 165 command! Black :py3 Black()
 166 command! BlackUpgrade :py3 BlackUpgrade()
 167 command! BlackVersion :py3 BlackVersion()