scripts/gn_complete/complete.py

#!/usr/bin/env python3.8
# Copyright 2019 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.

'''
Prints completions for GN build target names.

Reads the GN build targets from //out/default/project_lite.json
(generated by prepare_targets.py) and prints completions that match the search
label.

This can be invoked to generate completions in any shell that supports
completions.

For a description of GN labels, run `fx gn help labels`.
'''

import argparse
import collections
import json
import os
import sys

FUCHSIA_DIR = os.environ['FUCHSIA_DIR']
DEFAULT_PROJECT_FILE = 'project_lite.json'

try:
  DEFAULT_OUT_DIR = os.environ['FUCHSIA_BUILD_DIR']
except KeyError:
  DEFAULT_OUT_DIR = os.path.join(FUCHSIA_DIR, 'out', 'default')

class Completions(object):
  '''Provides shell-like completions for partial GN labels.'''

  def __init__(self, prefix):
    self._prefix = prefix
    def factory():
      return collections.defaultdict(factory)
    self._children = factory()

  def prefix(self):
    return self._prefix

  def insert_label(self, label):
    '''Inserts the `label` into the `completions` dict.

    Args:
      label: A GN label that must start with `self.prefix()`.

    Raises:
      Exception: If `label` does not start with `self.prefix()`.
    '''
    if not label.startswith(self._prefix):
      raise Exception('label "{}" must have the prefix "{}"'.format(
          label, self._prefix))
    children = self._children
    for segment in split_label(label[len(self._prefix):]):
      children = children[segment]

  def list_completions(self):
    '''Yields partial GN label completions truncated at the first ambiguity.

    Eg. for a set of targets:
      ['foo/bar/baz:test',
       'foo/bar/bat/bar:test',
       'foo/bar/bat:bin',
       'foo/bar:test']

    the expected completions are:
      ['foo/bar/baz',
       'foo/bar/bat',
       'foo/bar:test']

    Args:
      completions: The dictionary of completions.

    Yields:
      Partial GN labels that should be suggested as completions.
    '''
    if not self._children:
      return

    label = ''
    children = self._children

    # Descend down the label segments until there is a choice
    # to make. These are the most useful completions.
    while len(children) == 1:
      segment, children = next(iter(children.items()))
      label += segment

    if not children:
      # There is only a single match to suggest.
      yield self._prefix + label

    for segment, _ in children.items():
      yield self._prefix + label + segment


def get_project_subdir():
  '''Returns the current working subdirectory of FUCHSIA_DIR.

  Returns:
    A subdirectory without a leading '/'.

  Raises:
    Exception: if the current working directory is not a subdirectory
        of FUCHSIA_DIR.
  '''
  cwd = os.getcwd()
  if not cwd.startswith(FUCHSIA_DIR):
    raise Exception('current working directory must be in {}'.format(
        FUCHSIA_DIR))
  subdir = cwd[len(FUCHSIA_DIR):]
  if subdir.startswith('/'):
    subdir = subdir[1:]
  return subdir


def list_matching_targets(project_file, prefix):
  '''Yields GN targets in `project_file` that match `prefix`.

  Args:
    project_file: A JSON file which contains a list of GN labels under the
        'targets' key in the root object.
    prefix: The GN label prefix against which to match the list of targets.

  Returns:
    A generator that yields GN targets that match the GN label `prefix`.
  '''
  for target in json.load(project_file)['targets']:
    if target.startswith(prefix):
      yield target


def is_label_relative(label):
  '''Returns True if `label` is relative (starts with '//').'''
  return not label.startswith('//')


def concat_labels(base, leaf):
  '''Concatenates two GN labels, removing extra '/' characters if needed.

  Args:
    base: A relative or absolute GN label which does not include a target
          (:foo).
    leaf: A relative GN label.

  Returns:
    A GN label that is the concatenation of `base` and `leaf`. The label is
    relative or absolute depending on the `base` label.

  Raises:
    Exception: If the `leaf` label is absolute or `base` contains a target
               (:foo).
  '''
  if ':' in base:
    raise Exception('base label cannot contain a target: {}'.format(base))
  if leaf.startswith('/'):
    raise Exception('leaf label cannot be absolute: {}'.format(leaf))
  colon = leaf.startswith(':')
  if colon or not leaf or base.endswith('/'):
    return base + leaf
  return '{}/{}'.format(base, leaf)


def make_label_relative(base, absolute):
  '''Makes `absolute` a GN label relative to `base`.

  Args:
    base: An absolute GN label which does not include a target (:foo).
    absolute: An absolute GN label of which `base` is a prefix.

  Returns:
    A GN label relative to `base` that refers to the same label as `absolute`.
  '''
  relative = absolute[len(base):]
  if relative.startswith('/'):
    relative = relative[1:]
  return relative


def split_label(label):
  '''Splits a potentially partial GN label into segments.

  The segments are meant to be joined without a separator, which makes handling
  certain edge cases easier. Each delimiter ('/', ':') is included in the
  segment that follows it. Eg 'foo/bar:baz' => ['foo', '/bar', ':baz'].

  Absolute GN labels always start with the '//' segment.

  Toolchains are kept together with the target segment.

  Args:
    label: A potentially partial GN label, relative or absolute.

  Returns:
    A list of segments that can be joined without a separator to re-create
    the original label.
  '''
  if not label:
    return ['']

  toolchain = ''
  if label.endswith(')'):
    # This label contains a toolchain, which looks like :baz(//toolchain).
    toolchain_idx = label.rindex('(')
    toolchain = label[toolchain_idx:]
    label = label[0:toolchain_idx]

  leaf = ''
  leaf_idx = label.find(':')
  if leaf_idx >= 0:
    leaf = label[leaf_idx:]
    label = label[0:leaf_idx]
  segments = []
  if label.startswith('//'):
    segments = ['//']
    label = label[2:]
  while True:
    # Always search after the last /.
    idx = label.find('/', 1)
    if idx > 0:
      segments.append(label[0:idx])
      label = label[idx:]
    else:
      if label:
        segments.append(label)
      break
  if leaf:
    segments.append(leaf)
  if toolchain:
    if segments:
      segments[-1] += toolchain
    else:
      segments = [toolchain]
  return segments


def main(args):
  base_label = ''
  prefix = args.search_label
  relative = is_label_relative(args.search_label)
  if relative:
    base_label = '//{}'.format(get_project_subdir())
    prefix = concat_labels(base_label, args.search_label)

  eprint('Base:   {}'.format(base_label))
  eprint('Search: {}'.format(args.search_label))
  eprint('Prefix: {}'.format(prefix))

  completions = Completions(args.search_label)
  with open(os.path.join(args.build_dir, args.project_file)) as fin:
    for target in list_matching_targets(fin, prefix):
      if relative:
        target = make_label_relative(base_label, target)
      completions.insert_label(target)

  for completion in completions.list_completions():
    print(completion)


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


if __name__ == '__main__':
  p = argparse.ArgumentParser(description=__doc__)
  p.add_argument('--build_dir', default=DEFAULT_OUT_DIR)
  p.add_argument('--project_file', default=DEFAULT_PROJECT_FILE)
  p.add_argument('search_label', nargs='?', default='')
  try:
    main(p.parse_args())
  except Exception as e:
    eprint('error: {}'.format(e))
    sys.exit(1)