X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/52384bf0a3fe4b0d0a11daad6ef90b8cd948a773..e974fc3c52959e9f01bf62bdcd0d8c100ce78985:/autoload/black.vim

diff --git a/autoload/black.vim b/autoload/black.vim
index 29d8f2b..051fea0 100644
--- a/autoload/black.vim
+++ b/autoload/black.vim
@@ -3,8 +3,13 @@ import collections
 import os
 import sys
 import vim
-from distutils.util import strtobool
 
+def strtobool(text):
+  if text.lower() in ['y', 'yes', 't', 'true', 'on', '1']:
+    return True
+  if text.lower() in ['n', 'no', 'f', 'false', 'off', '0']:
+    return False
+  raise ValueError(f"{text} is not convertible to boolean")
 
 class Flag(collections.namedtuple("FlagBase", "name, cast")):
   @property
@@ -24,10 +29,12 @@ FLAGS = [
   Flag(name="fast", cast=strtobool),
   Flag(name="skip_string_normalization", cast=strtobool),
   Flag(name="quiet", cast=strtobool),
+  Flag(name="skip_magic_trailing_comma", cast=strtobool),
+  Flag(name="preview", cast=strtobool),
 ]
 
 
-def _get_python_binary(exec_prefix):
+def _get_python_binary(exec_prefix, pyver):
   try:
     default = vim.eval("g:pymode_python").strip()
   except vim.error:
@@ -36,7 +43,15 @@ def _get_python_binary(exec_prefix):
     return default
   if sys.platform[:3] == "win":
     return exec_prefix / 'python.exe'
-  return exec_prefix / 'bin' / 'python3'
+  bin_path = exec_prefix / "bin"
+  exec_path = (bin_path / f"python{pyver[0]}.{pyver[1]}").resolve()
+  if exec_path.exists():
+    return exec_path
+  # It is possible that some environments may only have python3
+  exec_path = (bin_path / f"python3").resolve()
+  if exec_path.exists():
+    return exec_path
+  raise ValueError("python executable not found")
 
 def _get_pip(venv_path):
   if sys.platform[:3] == "win":
@@ -49,9 +64,19 @@ def _get_virtualenv_site_packages(venv_path, pyver):
   return venv_path / 'lib' / f'python{pyver[0]}.{pyver[1]}' / 'site-packages'
 
 def _initialize_black_env(upgrade=False):
+  if vim.eval("g:black_use_virtualenv ? 'true' : 'false'") == "false":
+    if upgrade:
+      print("Upgrade disabled due to g:black_use_virtualenv being disabled.")
+      print("Either use your system package manager (or pip) to upgrade black separately,")
+      print("or modify your vimrc to have 'let g:black_use_virtualenv = 1'.")
+      return False
+    else:
+      # Nothing needed to be done.
+      return True
+
   pyver = sys.version_info[:3]
-  if pyver < (3, 6, 2):
-    print("Sorry, Black requires Python 3.6.2+ to run.")
+  if pyver < (3, 8):
+    print("Sorry, Black requires Python 3.8+ to run.")
     return False
 
   from pathlib import Path
@@ -65,7 +90,7 @@ def _initialize_black_env(upgrade=False):
     _executable = sys.executable
     _base_executable = getattr(sys, "_base_executable", _executable)
     try:
-      executable = str(_get_python_binary(Path(sys.exec_prefix)))
+      executable = str(_get_python_binary(Path(sys.exec_prefix), pyver))
       sys.executable = executable
       sys._base_executable = executable
       print(f'Creating a virtualenv in {virtualenv_path}...')
@@ -98,13 +123,49 @@ if _initialize_black_env():
   import black
   import time
 
-def Black():
+def get_target_version(tv):
+  if isinstance(tv, black.TargetVersion):
+    return tv
+  ret = None
+  try:
+    ret = black.TargetVersion[tv.upper()]
+  except KeyError:
+    print(f"WARNING: Target version {tv!r} not recognized by Black, using default target")
+  return ret
+
+def Black(**kwargs):
+  """
+  kwargs allows you to override ``target_versions`` argument of
+  ``black.FileMode``.
+
+  ``target_version`` needs to be cleaned because ``black.FileMode``
+  expects the ``target_versions`` argument to be a set of TargetVersion enums.
+
+  Allow kwargs["target_version"] to be a string to allow
+  to type it more quickly.
+
+  Using also target_version instead of target_versions to remain
+  consistent to Black's documentation of the structure of pyproject.toml.
+  """
   start = time.time()
   configs = get_configs()
+
+  black_kwargs = {}
+  if "target_version" in kwargs:
+    target_version = kwargs["target_version"]
+
+    if not isinstance(target_version, (list, set)):
+      target_version = [target_version]
+    target_version = set(filter(lambda x: x, map(lambda tv: get_target_version(tv), target_version)))
+    black_kwargs["target_versions"] = target_version
+
   mode = black.FileMode(
     line_length=configs["line_length"],
     string_normalization=not configs["skip_string_normalization"],
     is_pyi=vim.current.buffer.name.endswith('.pyi'),
+    magic_trailing_comma=not configs["skip_magic_trailing_comma"],
+    preview=configs["preview"],
+    **black_kwargs,
   )
   quiet = configs["quiet"]
 
@@ -117,9 +178,9 @@ def Black():
     )
   except black.NothingChanged:
     if not quiet:
-      print(f'Already well formatted, good job. (took {time.time() - start:.4f}s)')
+      print(f'Black: already well formatted, good job. (took {time.time() - start:.4f}s)')
   except Exception as exc:
-    print(exc)
+    print(f'Black: {exc}')
   else:
     current_buffer = vim.current.window.buffer
     cursors = []
@@ -136,7 +197,7 @@ def Black():
       except vim.error:
         window.cursor = (len(window.buffer), 0)
     if not quiet:
-      print(f'Reformatted in {time.time() - start:.4f}s.')
+      print(f'Black: reformatted in {time.time() - start:.4f}s.')
 
 def get_configs():
   filename = vim.eval("@%")
@@ -160,8 +221,17 @@ def BlackVersion():
 
 EndPython3
 
-function black#Black()
-  :py3 Black()
+function black#Black(...)
+    let kwargs = {}
+    for arg in a:000
+        let arg_list = split(arg, '=')
+        let kwargs[arg_list[0]] = arg_list[1]
+    endfor
+python3 << EOF
+import vim
+kwargs = vim.eval("kwargs")
+EOF
+  :py3 Black(**kwargs)
 endfunction
 
 function black#BlackUpgrade()