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     func! __BLACK_MISSING()
  18         echo "The black.vim plugin requires vim7.0+ with Python 3.6 support."
  19     endfunc
  20     command! Black :call __BLACK_MISSING()
  21     command! BlackUpgrade :call __BLACK_MISSING()
  22     command! BlackVersion :call __BLACK_MISSING()
  23     finish
  24 endif
  25
  26 if exists("g:load_black")
  27    finish
  28 endif
  29
  30 let g:load_black = "py1.0"
  31 if !exists("g:black_virtualenv")
  32   if has("nvim")
  33     let g:black_virtualenv = "~/.local/share/nvim/black"
  34   else
  35     let g:black_virtualenv = "~/.vim/black"
  36   endif
  37 endif
  38 if !exists("g:black_fast")
  39   let g:black_fast = 0
  40 endif
  41 if !exists("g:black_linelength")
  42   let g:black_linelength = 88
  43 endif
  44 if !exists("g:black_string_normalization")
  45   if exists("g:black_skip_string_normalization")
  46     let g:black_string_normalization = !g:black_skip_string_normalization
  47   else
  48     let g:black_string_normalization = 1
  49   endif
  50 endif
  51 if !exists("g:black_quiet")
  52   let g:black_quiet = 0
  53 endif
  54
  55 python3 << EndPython3
  56 import collections
  57 import os
  58 import sys
  59 import vim
  60 from distutils.util import strtobool
  61
  62
  63 class Flag(collections.namedtuple("FlagBase", "name, cast")):
  64   @property
  65   def var_name(self):
  66     return self.name.replace("-", "_")
  67
  68   @property
  69   def vim_rc_name(self):
  70     name = self.var_name
  71     if name == "line_length":
  72       name = name.replace("_", "")
  73     return "g:black_" + name
  74
  75
  76 FLAGS = [
  77   Flag(name="line_length", cast=int),
  78   Flag(name="fast", cast=strtobool),
  79   Flag(name="string_normalization", cast=strtobool),
  80   Flag(name="quiet", cast=strtobool),
  81 ]
  82
  83
  84 def _get_python_binary(exec_prefix):
  85   try:
  86     default = vim.eval("g:pymode_python").strip()
  87   except vim.error:
  88     default = ""
  89   if default and os.path.exists(default):
  90     return default
  91   if sys.platform[:3] == "win":
  92     return exec_prefix / 'python.exe'
  93   return exec_prefix / 'bin' / 'python3'
  94
  95 def _get_pip(venv_path):
  96   if sys.platform[:3] == "win":
  97     return venv_path / 'Scripts' / 'pip.exe'
  98   return venv_path / 'bin' / 'pip'
  99
 100 def _get_virtualenv_site_packages(venv_path, pyver):
 101   if sys.platform[:3] == "win":
 102     return venv_path / 'Lib' / 'site-packages'
 103   return venv_path / 'lib' / f'python{pyver[0]}.{pyver[1]}' / 'site-packages'
 104
 105 def _initialize_black_env(upgrade=False):
 106   pyver = sys.version_info[:2]
 107   if pyver < (3, 6):
 108     print("Sorry, Black requires Python 3.6+ to run.")
 109     return False
 110
 111   from pathlib import Path
 112   import subprocess
 113   import venv
 114   virtualenv_path = Path(vim.eval("g:black_virtualenv")).expanduser()
 115   virtualenv_site_packages = str(_get_virtualenv_site_packages(virtualenv_path, pyver))
 116   first_install = False
 117   if not virtualenv_path.is_dir():
 118     print('Please wait, one time setup for Black.')
 119     _executable = sys.executable
 120     _base_executable = getattr(sys, "_base_executable", _executable)
 121     try:
 122       executable = str(_get_python_binary(Path(sys.exec_prefix)))
 123       sys.executable = executable
 124       sys._base_executable = executable
 125       print(f'Creating a virtualenv in {virtualenv_path}...')
 126       print('(this path can be customized in .vimrc by setting g:black_virtualenv)')
 127       venv.create(virtualenv_path, with_pip=True)
 128     except Exception:
 129       print('Encountered exception while creating virtualenv (see traceback below).')
 130       print(f'Removing {virtualenv_path}...')
 131       import shutil
 132       shutil.rmtree(virtualenv_path)
 133       raise
 134     finally:
 135       sys.executable = _executable
 136       sys._base_executable = _base_executable
 137     first_install = True
 138   if first_install:
 139     print('Installing Black with pip...')
 140   if upgrade:
 141     print('Upgrading Black with pip...')
 142   if first_install or upgrade:
 143     subprocess.run([str(_get_pip(virtualenv_path)), 'install', '-U', 'black'], stdout=subprocess.PIPE)
 144     print('DONE! You are all set, thanks for waiting ✨ 🍰 ✨')
 145   if first_install:
 146     print('Pro-tip: to upgrade Black in the future, use the :BlackUpgrade command and restart Vim.\n')
 147   if virtualenv_site_packages not in sys.path:
 148     sys.path.insert(0, virtualenv_site_packages)
 149   return True
 150
 151 if _initialize_black_env():
 152   import black
 153   import time
 154
 155 def Black():
 156   start = time.time()
 157   configs = get_configs()
 158   mode = black.FileMode(
 159     line_length=configs["line_length"],
 160     string_normalization=configs["string_normalization"],
 161     is_pyi=vim.current.buffer.name.endswith('.pyi'),
 162   )
 163   quiet = configs["quiet"]
 164
 165   buffer_str = '\n'.join(vim.current.buffer) + '\n'
 166   try:
 167     new_buffer_str = black.format_file_contents(
 168       buffer_str,
 169       fast=configs["fast"],
 170       mode=mode,
 171     )
 172   except black.NothingChanged:
 173     if not quiet:
 174       print(f'Already well formatted, good job. (took {time.time() - start:.4f}s)')
 175   except Exception as exc:
 176     print(exc)
 177   else:
 178     current_buffer = vim.current.window.buffer
 179     cursors = []
 180     for i, tabpage in enumerate(vim.tabpages):
 181       if tabpage.valid:
 182         for j, window in enumerate(tabpage.windows):
 183           if window.valid and window.buffer == current_buffer:
 184             cursors.append((i, j, window.cursor))
 185     vim.current.buffer[:] = new_buffer_str.split('\n')[:-1]
 186     for i, j, cursor in cursors:
 187       window = vim.tabpages[i].windows[j]
 188       try:
 189         window.cursor = cursor
 190       except vim.error:
 191         window.cursor = (len(window.buffer), 0)
 192     if not quiet:
 193       print(f'Reformatted in {time.time() - start:.4f}s.')
 194
 195 def get_configs():
 196   path_pyproject_toml = black.find_pyproject_toml(vim.eval("fnamemodify(getcwd(), ':t')"))
 197   if path_pyproject_toml:
 198     toml_config = black.parse_pyproject_toml(path_pyproject_toml)
 199   else:
 200     toml_config = {}
 201
 202   return {
 203     flag.var_name: flag.cast(toml_config.get(flag.name, vim.eval(flag.vim_rc_name)))
 204     for flag in FLAGS
 205   }
 206
 207
 208 def BlackUpgrade():
 209   _initialize_black_env(upgrade=True)
 210
 211 def BlackVersion():
 212   print(f'Black, version {black.__version__} on Python {sys.version}.')
 213
 214 EndPython3
 215
 216 command! Black :py3 Black()
 217 command! BlackUpgrade :py3 BlackUpgrade()
 218 command! BlackVersion :py3 BlackVersion()