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 "The black.vim plugin 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 collections
  45 import os
  46 import sys
  47 import vim
  48
  49
  50 class Flag(collections.namedtuple("FlagBase", "name, cast")):
  51   @property
  52   def var_name(self):
  53     return self.name.replace("-", "_")
  54
  55   @property
  56   def vim_rc_name(self):
  57     name = self.var_name
  58     if name == "line_length":
  59       name = name.replace("_", "")
  60     if name == "string_normalization":
  61       name = "skip_" + name
  62     return "g:black_" + name
  63
  64
  65 FLAGS = [
  66   Flag(name="line_length", cast=int),
  67   Flag(name="fast", cast=bool),
  68   Flag(name="string_normalization", cast=bool),
  69 ]
  70
  71
  72 def _get_python_binary(exec_prefix):
  73   try:
  74     default = vim.eval("g:pymode_python").strip()
  75   except vim.error:
  76     default = ""
  77   if default and os.path.exists(default):
  78     return default
  79   if sys.platform[:3] == "win":
  80     return exec_prefix / 'python.exe'
  81   return exec_prefix / 'bin' / 'python3'
  82
  83 def _get_pip(venv_path):
  84   if sys.platform[:3] == "win":
  85     return venv_path / 'Scripts' / 'pip.exe'
  86   return venv_path / 'bin' / 'pip'
  87
  88 def _get_virtualenv_site_packages(venv_path, pyver):
  89   if sys.platform[:3] == "win":
  90     return venv_path / 'Lib' / 'site-packages'
  91   return venv_path / 'lib' / f'python{pyver[0]}.{pyver[1]}' / 'site-packages'
  92
  93 def _initialize_black_env(upgrade=False):
  94   pyver = sys.version_info[:2]
  95   if pyver < (3, 6):
  96     print("Sorry, Black requires Python 3.6+ to run.")
  97     return False
  98
  99   from pathlib import Path
 100   import subprocess
 101   import venv
 102   virtualenv_path = Path(vim.eval("g:black_virtualenv")).expanduser()
 103   virtualenv_site_packages = str(_get_virtualenv_site_packages(virtualenv_path, pyver))
 104   first_install = False
 105   if not virtualenv_path.is_dir():
 106     print('Please wait, one time setup for Black.')
 107     _executable = sys.executable
 108     try:
 109       sys.executable = str(_get_python_binary(Path(sys.exec_prefix)))
 110       print(f'Creating a virtualenv in {virtualenv_path}...')
 111       print('(this path can be customized in .vimrc by setting g:black_virtualenv)')
 112       venv.create(virtualenv_path, with_pip=True)
 113     finally:
 114       sys.executable = _executable
 115     first_install = True
 116   if first_install:
 117     print('Installing Black with pip...')
 118   if upgrade:
 119     print('Upgrading Black with pip...')
 120   if first_install or upgrade:
 121     subprocess.run([str(_get_pip(virtualenv_path)), 'install', '-U', 'black'], stdout=subprocess.PIPE)
 122     print('DONE! You are all set, thanks for waiting ✨ 🍰 ✨')
 123   if first_install:
 124     print('Pro-tip: to upgrade Black in the future, use the :BlackUpgrade command and restart Vim.\n')
 125   if virtualenv_site_packages not in sys.path:
 126     sys.path.append(virtualenv_site_packages)
 127   return True
 128
 129 if _initialize_black_env():
 130   import black
 131   import time
 132
 133 def Black():
 134   start = time.time()
 135   configs = get_configs()
 136   mode = black.FileMode(
 137     line_length=configs["line_length"],
 138     string_normalization=configs["string_normalization"],
 139     is_pyi=vim.current.buffer.name.endswith('.pyi'),
 140   )
 141
 142   buffer_str = '\n'.join(vim.current.buffer) + '\n'
 143   try:
 144     new_buffer_str = black.format_file_contents(
 145       buffer_str,
 146       fast=configs["fast"],
 147       mode=mode,
 148     )
 149   except black.NothingChanged:
 150     print(f'Already well formatted, good job. (took {time.time() - start:.4f}s)')
 151   except Exception as exc:
 152     print(exc)
 153   else:
 154     current_buffer = vim.current.window.buffer
 155     cursors = []
 156     for i, tabpage in enumerate(vim.tabpages):
 157       if tabpage.valid:
 158         for j, window in enumerate(tabpage.windows):
 159           if window.valid and window.buffer == current_buffer:
 160             cursors.append((i, j, window.cursor))
 161     vim.current.buffer[:] = new_buffer_str.split('\n')[:-1]
 162     for i, j, cursor in cursors:
 163       window = vim.tabpages[i].windows[j]
 164       try:
 165         window.cursor = cursor
 166       except vim.error:
 167         window.cursor = (len(window.buffer), 0)
 168     print(f'Reformatted in {time.time() - start:.4f}s.')
 169
 170 def get_configs():
 171   path_pyproject_toml = black.find_pyproject_toml(vim.eval("fnamemodify(getcwd(), ':t')"))
 172   if path_pyproject_toml:
 173     toml_config = black.parse_pyproject_toml(path_pyproject_toml)
 174   else:
 175     toml_config = {}
 176
 177   return {
 178     flag.var_name: toml_config.get(flag.name, flag.cast(vim.eval(flag.vim_rc_name)))
 179     for flag in FLAGS
 180   }
 181
 182
 183 def BlackUpgrade():
 184   _initialize_black_env(upgrade=True)
 185
 186 def BlackVersion():
 187   print(f'Black, version {black.__version__} on Python {sys.version}.')
 188
 189 endpython3
 190
 191 command! Black :py3 Black()
 192 command! BlackUpgrade :py3 BlackUpgrade()
 193 command! BlackVersion :py3 BlackVersion()