#! /usr/bin/env python

# Magic line for bash (blame obscure shells):
#
#    complete -o bashdefault -o default -C /path/to/completepytest.py py.test

import sys

complete_with_test_prefix = True

linestarts = [("def ", "test_"), ("class ", "Test")]

command, wordstart, previousword = sys.argv[1:4]
if wordstart == '-k':
    wordstart = ''
    addspace = '-k '
else:
    if previousword != '-k':
        sys.exit(0)    # use default completion
    addspace = ''

import os
line = os.environ['COMP_LINE']

filenames = [arg for arg in line.split()
             if (arg.endswith('.py') or arg.endswith('.pl')) and
                os.path.basename(arg).startswith('test_')]
place = wordstart.find('.')
if '.' in wordstart:
    addspace = wordstart[:place + 1]
    wordstart = wordstart[place + 1:]

def gen_lines(arg):
    try:
        f = open(arg, 'r')
        lines = f.readlines()
        f.close()
    except (OSError, IOError):
        return
    for line in lines:
        yield line.lstrip()

def complete_python(arg, filename):
    for line in gen_lines(arg):
        for (keyword, prefix) in linestarts:
            if not line.startswith(keyword + prefix):
                continue
            testname = line[len(keyword):]
            i = testname.find('(')
            if i < 0:
                i = testname.find(':')
                if i < 0:
                    continue
            if complete_with_test_prefix:
                testname = testname[:i]
                if testname.startswith(wordstart):
                    print addspace + testname
            else:
                testname = testname[len(prefix):i]
                if testname.startswith(wordstart):
                    print addspace + testname
            break
        else:
            if line.startswith('# =='):
                # follow comments of the form:
                # ===> relative_file_path.py
                parts = line.split(None, 2)
                if len(parts) == 3:
                    fn = parts[2].strip()
                    fn = os.path.join(os.path.dirname(arg), fn)
                    fn = os.path.abspath(fn)
                    if fn not in filenames:
                        filenames.append(fn)

def complete_prolog(arg, filenames):
    for line in gen_lines(arg):
        if line.startswith("test_"):
            i = line.find(' :-')
            if i < 0:
                continue
            testname = line[:i]
            if testname.startswith(wordstart):
                print addspace + testname

for arg in filenames:
    if arg.endswith(".py"):
        complete_python(arg, filenames)
    else:
        assert arg.endswith(".pl")
        complete_prolog(arg, filenames)