#!/usr/bin/pypy3

import argparse
import itertools
import os
import re
import py_compile
import subprocess
import sys


def abort(message):
    print(message, file=sys.stderr)
    sys.exit(1)


def package_modules(package):
    """Iterate through all python modules in an installed Debian package"""
    p = subprocess.Popen(('dpkg', '-L', package), stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    files, stderr = p.communicate()
    if p.returncode != 0:
        abort('Unable to list files in %s. Is it installed?' % package)

    for fn in files.decode('utf8').splitlines():
        if fn.endswith('.py'):
            if fn.startswith('/usr/share/doc/'):
                continue
            yield fn


def find_modules(root):
    """Iterate through all python modules in directory tree root"""
    if os.path.isfile(root):
        yield root
        return

    for dirpath, dirnames, filenames in os.walk(root):
        for fn in filenames:
            if fn.endswith('.py'):
                yield os.path.join(dirpath, fn)


# BCEP support, from py3compile
def get_exclude_patterns_from_dir(name='/usr/share/python3/bcep/'):
    """Return patterns for files that shouldn't be bytecompiled."""
    if not os.path.isdir(name):
        return

    for fn in os.listdir(name):
        if fn.startswith('.'):
            continue
        with open(os.path.join(name, fn), 'r', encoding='utf-8') as lines:
            for line in lines:
                if line.startswith('#'):
                    continue
                line = line.rstrip('\n')

                type_, vrange, dname = line.split('|', 2)
                if not within_vrange(vrange):
                    continue

                try:
                    dname, pattern = dname.split('|', 1)
                except ValueError:
                    pattern = '.*'

                if type_ != 'file' and not dname.endswith('/'):
                    dname += '/'

                if pattern:
                    try:
                        pattern = re.compile(pattern)
                    except re.error:
                        print('skipping invalid pattern in file {}, line: {}'
                              .format(fn, line), file=sys.stderr)
                        continue

                yield type_, dname, pattern


def within_vrange(vrange):
    """Is our interpreter's cpython compatibility level within vrange?"""
    min_, sep, max_ = vrange.partition('-')
    if min_:
        min_ = tuple(map(int, min_.split('.')))
    if max_:
        max_ = tuple(map(int, max_.split('.')))
    version = sys.version_info[:2]

    if not min_ and not max_:
        return True

    if not sep:
        return version == min_

    if max_ and version >= max_:
        return False

    if min_ and min_ > version:
        return False

    return True


def main():
    parser = argparse.ArgumentParser(
        description='Byte-compile Python source files in a package, for PyPy')
    parser.add_argument('-p', '--package', metavar='PACKAGE',
                        action='append', default=[],
                        help='Debian package to byte-compile '
                             '(may be specified multiple times)')
    parser.add_argument('directory', nargs='*',
                        help='Directory tree (or file) to byte-compile')
    parser.add_argument('-X', '--exclude', metavar='REGEXPR',
                        action='append', default=[], type=re.compile,
                        help='Exclude items that match given REGEXPR '
                             '(may be specified multiple times)')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Be more verbose')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Be quiet')
    parser.add_argument('-V', metavar='VRANGE', dest='vrange',
                        help='Only byte-compile modules if the cPython '
                             'compatibility version falls within the specified '
                             'range.')
    parser.add_argument('-O', action='store_true', dest='pyo',
                        help=argparse.SUPPRESS)
    parser.add_argument('--bcep-path', default='/usr/share/python3/bcep/',
                        help=argparse.SUPPRESS)
    args = parser.parse_args()

    if not (args.package or args.directory):
        parser.error('Either a package or a directory must be specified')
    if args.quiet and args.verbose:
        parser.error('--quiet and --verbose cannot both be specified')

    if args.vrange and not within_vrange(args.vrange):
        return

    if args.pyo:
        print('-O is ignored in pypy3compile', file=sys.stderr)

    modules_p = set(itertools.chain(*(
        package_modules(package) for package in args.package)))
    modules_d = set(itertools.chain(*(
        find_modules(dir_) for dir_ in args.directory)))

    if args.package and args.directory:
        modules = modules_d & modules_p
    else:
        modules = modules_d | modules_p

    modules = filter(lambda module: not any(pattern.match(module)
                                            for pattern in args.exclude),
                     modules)

    bcep_excludes = list(get_exclude_patterns_from_dir(args.bcep_path))

    for module in modules:
        skip = False
        for pattern in args.exclude:
            if pattern.match(module):
                skip = True
        for type_, dname, pattern in bcep_excludes:
            if not module.startswith(dname):
                continue
            if type_ == 'dir':
                skip = True
            if type_ == 'file' and module == dname:
                skip = True
            if type_ == 're':
                relative_path = module[len(dname):]
                if pattern.match(module) or pattern.match(relative_path):
                    skip = True
        if skip:
            continue
        if args.verbose:
            print('Byte-compiling %s' % module)
        try:
            py_compile.compile(module, doraise=True)
        except py_compile.PyCompileError as e:
            if not args.quiet:
                print('Failed to byte-compile %s: %s' % (module, e.msg),
                      file=sys.stderr)


if __name__ == '__main__':
    main()

# vim: ft=python
