#!/usr/bin/pypy3

import argparse
import collections
import itertools
import os
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)


def clean_modules(modules, verbose, tag):
    """Remove all tag.pyc files for every module specified"""
    clean = collections.defaultdict(list)
    for module in modules:
        dir_, basename = os.path.split(module)
        clean[dir_].append(os.path.splitext(basename)[0])

    for dir_, basenames in clean.items():
        pycache = os.path.join(dir_, '__pycache__')
        if not os.path.exists(pycache):
            continue

        empty = True
        for fn in os.listdir(pycache):
            parts = fn.rsplit('.', 2)
            if parts[0] in basenames and parts[1] == tag and parts[2] == 'pyc':
                if verbose:
                    print('Removing %s' % os.path.join(pycache, fn))
                os.unlink(os.path.join(pycache, fn))
            else:
                empty = False

        if empty:
            if verbose:
                print('Pruning %s' % pycache)
            os.rmdir(pycache)


def clean_directories(directories, verbose, tag):
    """Remove the specified tag.pyc files."""
    suffix = '.{}.pyc'.format(tag)
    for root in directories:
        for dirpath, dirnames, filenames in os.walk(root):
            if dirpath.endswith('/__pycache__'):
                empty = True
                for file_ in filenames:
                    if file_.endswith(suffix):
                        if verbose:
                            print('Removing %s' % os.path.join(dirpath, file_))
                        os.unlink(os.path.join(dirpath, file_))
                    else:
                        empty = False

                if empty:
                    if verbose:
                        print('Pruning %s' % dirpath)
                    os.rmdir(dirpath)


def pypy3_tag(version):
    """Return the tag used by the given PyPy3 version"""
    parts = version.split('.')
    if len(parts) == 3:
        parts.pop()
    if len(parts) != 2:
        raise argparse.ArgumentTypeError('Version must be MAJOR.MINOR')
    parts = (int(parts[0]), int(parts[1]))
    if parts < (7, 3):
        return 'pypy3-{}{}'.format(*parts)
    # 7.3 started to use the cPython version:
    return 'pypy38'  # Compatible with 3.8


def main():
    parser = argparse.ArgumentParser(
        description='Remove byte-compiled files for a package')
    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('-v', '--verbose', action='store_true',
                        help='Be more verbose')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Be quiet')
    parser.add_argument('--version', metavar='VERSION',
                        default='{}.{}'.format(*sys.pypy_version_info[:2]),
                        dest='tag', type=pypy3_tag,
                        help='Clean up .pyc files left by the specified '
                             'pypy3 version')
    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')

    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
    elif args.package:
        modules = modules_p
    else:
        # Split files from directories, so that we can completely clean any
        # specified directories.
        modules = set()
        directories = set()
        for fn in args.directory:
            if os.path.isfile(fn) and fn.endswith('.py'):
                modules.add(fn)
            else:
                directories.add(fn)
        clean_directories(directories, args.verbose, args.tag)

    clean_modules(modules, args.verbose, args.tag)


if __name__ == '__main__':
    main()

# vim: ft=python
