find_short_prefix.py

#! /usr/bin/env python3
""" Find conda packages which have a prefix. """

import argparse
import json
import os
import urllib
import tarfile
import xmlrpc.client as xmlrpclib

import conda.api as api

try:
    from packaging.version import parse as parse_version
except ImportError:
    from pip._vendor.packaging.version import parse as parse_version


def find_latest_versions(index, package_name):
    """ Return the latest version and packages from a conda channel index. """
    valid = [v for v in index.values() if v['name'] == package_name]
    versions = [parse_version(v['version']) for v in valid]
    latest_ver = max(versions)
    base_ver = latest_ver.base_version
    entries = [v for v in valid if v['version'] == base_ver]
    return latest_ver, entries


def parse_arguments():
    """ Parse command line arguments. """
    parser = argparse.ArgumentParser(
        description="Find conda packages which use a prefix")
    parser.add_argument(
        'packages', nargs='*',
        help=('Name of packages to check, leave blank to check all packages '
              'on the channel'))
    parser.add_argument(
        '--skip', '-s', action='store', help=(
            'file containing list of packages to skip when checking for '
            'prefixes'))
    parser.add_argument(
        '--verb', '-v', action='store_true', help='verbose output')
    parser.add_argument(
        '--channel', '-c', action='store', default='conda-forge',
        help='Conda channel to check.  Default is conda-forge')
    parser.add_argument(
        '--json', action='store', help='Save outdated packages to json file.')
    parser.add_argument(
        '--directory', '-d', action='store', default=os.path.join(os.getcwd(), 'pkg_cache'), help='where to store packages')
    return parser.parse_args()


def find_prefix_packages(index, package_names, verbose, cache_dir):
    """ Return a list of packages which use a prefix. """
    client = xmlrpclib.ServerProxy('https://pypi.python.org/pypi')

    pkgs_with_prefix = []
    unknown = []
    no_prefix = []

    for package_name in sorted(package_names):
        latest_ver, entries = find_latest_versions(index, package_name)
        has_prefix = [None] # [e.get('has_prefix') for e in entries]

        if not entries:
            print(package_name + " : Missing any entries. Skipping...")
            continue

        if None in has_prefix:
            # Download if not in cache
            
            # sort entired by md5 so we try the same package each time
            entries = sorted(entries, key=lambda k: k['md5'])
           
            url = entries[0]['url']
            filename = os.path.join(cache_dir, url.split('/')[-1])
            if not os.path.exists(filename):
                print("Downloading:", filename)
                response = urllib.request.urlopen(url)
                with open(filename, 'wb') as f:
                    f.write(response.read())

            tf = tarfile.open(filename)
            try:
                uses_prefix = b' binary ' in tf.extractfile(tf.getmember('info/has_prefix')).read()
                #print("!!!", package_name, 'good')
            except KeyError:
                uses_prefix = False
                #print("!!!", package_name, 'bad')

        else:
            uses_prefix = any(has_prefix)

        if uses_prefix is None:
            print(package_name, "UNKNOWN if uses a prefix")
            unknown.append(package_name)

        elif uses_prefix:
            print(package_name, "use a prefix")
            pkgs_with_prefix.append(package_name)

        elif uses_prefix is False:
            no_prefix.append(package_name)
            if verbose:
                print(package_name, "does NOT use a prefix")

    print("Unknown:", len(unknown))
    print("Prefix:", len(pkgs_with_prefix))
    print("No Prefix:", len(no_prefix))

    return pkgs_with_prefix


def main():
    args = parse_arguments()
    
    # create somewhere to store downloaded packages.
    if not os.path.exists(args.directory):
        os.makedirs(args.directory)

    # determine package names to check
    index = api.get_index(
        channel_urls=[args.channel], prepend=False, use_cache=False)
    package_names = set(args.packages)
    if len(package_names) == 0:  # no package names given on command line
        package_names = {v['name'] for k, v in index.items()}

    # remove skipped packages
    if args.skip is not None:
        with open(args.skip) as f:
            pkgs_to_skip = [line.strip() for line in f]
        package_names = [p for p in package_names if p not in pkgs_to_skip]

    outdated_packages = find_prefix_packages(index, package_names, args.verb, args.directory)

    # save outdated_packages to json formatted file is specified
    if args.json is not None:
        with open(args.json, 'w') as f:
            json.dump(outdated_packages, f)


if __name__ == "__main__":
    main()