fd.py

#!/usr/bin/env python
# Copyright 2017 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""fd.py is a fascinating directory changer to save your time in typing.

Use shell function "fd()" to enable autocompletion.
Use "fd.py" directly for without autocompletion.

See examples by
$ fd.py --help
"""

from __future__ import print_function

import argparse
import os
import pickle
import sys
import termios
import tty

SEARCH_BASE = os.environ['FUCHSIA_DIR']  # or 'HOME'
TMP_BASE = '/tmp/'
DIRS_FILE = TMP_BASE + 'fd.txt'
PICKLE_FILE = TMP_BASE + 'fd.pickle'

EXCLUDE_DIRS = [
    '"*/.git"', './build', './buildtools', './out', './third_party',
    './zircon/build', './zircon/prebuilt', './cmake-build-debug', './zircon/third_party',
]


def eprint(*args, **kwargs):
  print(*args, file=sys.stderr, **kwargs)


class Trie(object):
  """Class Trie.
  """

  def __init__(self):
    self.name = ''  # != path up to here from the root. Key is a valid
    # complete one.
    self.vals = []
    self.kids = {}

  def __getitem__(self, name, idx=0):
    if self.name == name:
      return self.vals
    if idx == name.__len__() or name[idx] not in self.kids:
      return None
    return self.kids[name[idx]].__getitem__(name, idx + 1)

  def __setitem__(self, name, val, idx=0):
    if idx < name.__len__():
      self.kids.setdefault(name[idx], Trie()).__setitem__(name, val, idx + 1)
      return
    self.name = name
    self.vals.append(val)

  def __contains__(self, name):
    return self[name] is not None

  def walk(self):
    descendants = []
    if self.name:
      descendants.append(self.name)
    for k in self.kids:
      descendants.extend(self.kids[k].walk())
    return descendants

  def prefixed(self, name, idx=0):
    if idx < name.__len__():
      if name[idx] in self.kids:
        return self.kids[name[idx]].prefixed(name, idx + 1)
      return []

    return self.walk()


def build_trie():
  """build trie.

  Returns:
    Trie
  """

  def build_find_cmd():
    paths = []
    for path in EXCLUDE_DIRS:
      paths.append('{} {}'.format('-path', path))
    return (r'cd {}; find . \( {} \) -prune -o -type d -print > '
            '{}').format(SEARCH_BASE, ' -o '.join(paths), DIRS_FILE)

  cmd_str = build_find_cmd()
  os.system(cmd_str)

  t = Trie()
  with open(DIRS_FILE, 'r') as f:
    for line in f:
      line = line[2:][:-1]
      tokens = line.split('/')
      if tokens.__len__() == 0:
        continue
      target = tokens[-1]

      t[target] = line

  return t


def get_trie():
  """get_trie.

  Returns:
    trie
  """

  def save_pickle(obj):
    with open(PICKLE_FILE, 'wb+') as f:
      pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)

  def load_pickle():
    with open(PICKLE_FILE, 'rb') as f:
      return pickle.load(f)

  if os.path.exists(PICKLE_FILE):
    return load_pickle()

  t = build_trie()
  save_pickle(t)
  return t


def button(idx):
  """button maps idx to an ascii value.
  """
  ascii = 0
  if 0 <= idx <= 8:
    ascii = ord('1') + idx
  elif 9 <= idx <= 34:
    ascii = ord('a') + idx - 9
  elif 35 <= idx <= 60:
    ascii = ord('A') + idx - 35
  elif 61 <= idx <= 75:
    ascii = ord('!') + idx - 61
  return str(unichr(ascii))


def get_button():  # Unix way
  fd = sys.stdin.fileno()
  old_settings = termios.tcgetattr(fd)
  try:
    tty.setraw(sys.stdin.fileno())
    ch = sys.stdin.read(1)
  finally:
    termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
  return ch


def choose_options(t, key, choice):
  # Build options by the given key
  if key in t:
    options = t[key]
  else:
    prefixed_keys = t.prefixed(key)
    options = []
    for pk in prefixed_keys:
      options.extend(t[pk])

  options = sorted(options)
  if options.__len__() == 0:
    eprint('No such directory: {}'.format(key))
    return None
  elif options.__len__() == 1:
    return options[0]
  elif options.__len__() > 75:  # See def button() for the limit.
    eprint('Too many ({}) results for "{}". '
           'Refine your prefix or time to buy 4K '
           'monitor\n').format(options.__len__(), key)
    return None

  def list_choices(l):
    for i in range(l.__len__()):
      eprint('[{}] {}'.format(button(i), l[i]))
    eprint()

  choice_dic = {}
  for idx, val in enumerate(options):
    choice_dic[button(idx)] = val

  if choice is not None and choice not in choice_dic:
    # Invalid pre-choice
    eprint('Choice "{}" not available\n'.format(choice))

  if choice not in choice_dic:
    list_choices(options)
    choice = get_button()
    if choice not in choice_dic:
      return None

  return choice_dic[choice]


def main():

  def parse_cmdline():
    example_commands = """

[eg] # Use "fd" for autocompletion (See //scripts/fx-env.sh)
  $ fd ral        # change directory to an only option: ralink
  $ fd wlan       # shows all "wlan" directories and ask to choose
  $ fd wlan 3     # change directory matching to option 3 of "fd wlan"
  $ fd [TAB]      # Autocomplete subdirectories from the current directory
  $ fd //[TAB]    # Autocomplete subdirectories from ${FUCHSIA_DIR}
  $ fd --rebuild  # rebuilds the directory structure cache
"""
    p = argparse.ArgumentParser(
        description='A fascinating directory changer',
        epilog=example_commands,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    p.add_argument(
        '--rebuild', action='store_true', help='rebuild the directory DB')
    p.add_argument('--base', type=str, default=None)
    p.add_argument('target', nargs='?', default='')
    p.add_argument('choice', nargs='?', default=None)

    # Redirect help messages to stderr
    if len(sys.argv) == 2:
      if sys.argv[1] in ['-h', '--help']:
        eprint(p.format_help())
        print('.')  # Stay at the current directory
        sys.exit(0)

    return p.parse_args()

  def get_abs_path(relative_dir):
    if relative_dir is not None:
      return os.path.join(SEARCH_BASE, relative_dir)
    return os.getcwd()

  def derive_dest(target):
    if not target:
      # To test if this command was invoked just to rebuild
      return get_abs_path('.') if args.rebuild is False else os.getcwd()

    if target[:2] == '//':
      target = target[2:]

    candidate = target

    # Do not guess-work when the user specifies an option to intend to use.
    # Do guess work otherwise.
    if not args.choice:
      if os.path.exists(candidate):
        return candidate

      candidate = get_abs_path(target)
      if os.path.exists(candidate):
        return candidate

      candidate = os.path.abspath(target)
      if os.path.exists(candidate):
        return candidate

    t = get_trie()
    return get_abs_path(choose_options(t, target, args.choice))

  args = parse_cmdline()
  if args.base:
    global SEARCH_BASE
    SEARCH_BASE = args.base

  if args.rebuild:
    os.remove(PICKLE_FILE)

  dest = derive_dest(args.target)
  dest = os.path.normpath(dest)
  print(dest)


if __name__ == '__main__':
  try:
    main()
  except Exception as e:  # Catch all
    eprint(e.message, e.args)
    print('.')  # Stay at the current directory upon exception