Futureize script

#!/usr/bin/env python3

"""Usage:

* ./update.py folder  -> preview which files in folder need changes
* ./update.py folder file.py target  -> update file
* ./update.py folder file.py --preview  -> preview changes to that file
* ./update.py folder file.py target --test  -> update file with tests folder as root
* ./update.py folder file.py --test --preview  -> update file with tests folder as root
"""

import subprocess
import itertools
import argparse

from typing import List


SRC_ROOT = 'src/python/pants'
TEST_ROOT = 'tests/python/pants_test'


def main() -> None:
  parser = create_parser()
  args = parser.parse_args()
  base_root = TEST_ROOT if args.test else SRC_ROOT
  target_root = f'{base_root}/{args.folder}'
  if not args.file_target:
    check_what_needs_changes(f'{target_root}/*.py')
    return
  file_path = f'{target_root}/{args.file_target}'
  pants_target_path = f'{target_root}:{args.pants_target}'
  # execute
  if args.preview:
    preview_changes(file_path)
    return
  call_futurize(file_path)
  if not file_changed(file_path):
    return
  move_imports_below_future(file_path)
  update_build_dependencies(target_root, args.pants_target)
  call_pants_fmt(pants_target_path)
  call_pants_test(pants_target_path)


def create_parser() -> argparse.ArgumentParser:
  parser = argparse.ArgumentParser()
  parser.add_argument('folder')
  parser.add_argument('file_target', nargs='?', default=None)
  parser.add_argument('pants_target', nargs='?', default=None)
  parser.add_argument('-t', '--test', action='store_true')
  parser.add_argument('-p', '--preview', action='store_true')
  return parser


def check_what_needs_changes(folder: str) -> None:
  subprocess.run(
    f'build-support/pants_dev_deps.venv/bin/futurize --stage2 --no-diffs {folder}',
    shell=True
  )


def call_futurize(path: str) -> None:
  subprocess.run([
    'build-support/pants_dev_deps.venv/bin/futurize',
    '--write',
    '--nobackup',
    '--stage2',
    path
  ])


def preview_changes(path: str) -> None:
  subprocess.run([
    'build-support/pants_dev_deps.venv/bin/futurize',
    '--stage2',
    path
  ])


def file_changed(path: str) -> bool:
  git_changes = subprocess.run(['git', 'ls-files', '-m'],
                               stdout=subprocess.PIPE,
                               encoding='utf-8').stdout.strip()
  return path in git_changes


def move_imports_below_future(path: str) -> None:

    def read_lines_before_future_import(path: str) -> List[str]:
      with open(path, 'r') as f:
        return list(itertools.takewhile(lambda line: 'from __future__ import' not in line, f))

    # find line numbers
    header_lines = read_lines_before_future_import(path)
    new_imports_above_future = [i for i, line in enumerate(header_lines)
                                if '#' not in line and line != '\n']
    future_import_first_line = len(header_lines)
    future_import_last_line = len(header_lines) + 1

    # skip if not needed
    if not new_imports_above_future:
      return

    # get lines
    with open(path, 'r') as f:
      lines = list(f.readlines())

    # overwrite
    with open(path, 'w') as f:
      for index in new_imports_above_future:
        lines.insert(future_import_last_line + 1, lines[index])
      pop_index = future_import_first_line - 1
      for _ in new_imports_above_future:
        lines.pop(pop_index)
        pop_index -= 1
      f.writelines(lines)


def call_pants_fmt(pants_target_path: str) -> None:
  subprocess.run([
    './pants',
    'fmt',
    pants_target_path
  ])


def call_pants_test(pants_test_target_path: str) -> None:
  subprocess.run([
    './pants',
    'test',
    pants_test_target_path
  ])


def update_build_dependencies(folder: str, pants_target: str) -> None:
  build_file = f'{folder}/BUILD'
  with open(build_file, 'r') as f:
    lines = list(f.readlines())
  target_index = lines.index(f"  name = '{pants_target}',\n")
  for i, line in enumerate(lines[target_index:]):
    if 'dependencies = [' in line:
      lines.insert(target_index + i + 1, "    '3rdparty/python:future',\n")
      break
    if line == '\n':  # dependencies section doesn't exist for target
      lines.insert(target_index + 2, '  dependencies = [\n')
      lines.insert(target_index + 3, "    '3rdparty/python:future',\n")
      lines.insert(target_index + 4, '  ]\n')
  with open(build_file, 'w') as f:
    f.writelines(lines)


if __name__ == '__main__':
    main()