combine_arch_meta.py

#!/usr/bin/env python
"""
Make a combined arch-specific core package list from a set of json files with lists of
packages in conda environments. The json files can be generated by running `conda list --json`
within the environment one wants to reproduce.
"""

import json
import jinja2
import argparse
import re
import pathlib
import yaml
import logging


def parser():
    parser_ = argparse.ArgumentParser(description=__doc__)
    parser_.add_argument("--name", help="name of the package", required=True, default="ska3-core")
    parser_.add_argument("--version", help="new package version", required=True, default="")
    parser_.add_argument("--env", action="append", metavar="list.json",
                         help="environment file (produced with `conda list --json`)", default=[])
    parser_.add_argument("--subtract-env", metavar="list.json",
                         help="Exclude packages in given json file",
                         action="append", default=[])
    parser_.add_argument("--in", metavar="meta.yaml", dest='_in',
                         help="Only include packages in given meta.yaml")
    parser_.add_argument("--not-in", metavar="meta.yaml",
                         help="Exclude packages in given meta.yaml")
    parser_.add_argument("--exclude", help="Exclude specific packages",
                         metavar="PKG_NAME", action="append", default=[])
    parser_.add_argument("--include", help="Include specific packages",
                         metavar="PKG_NAME", action="append", default=[])
    parser_.add_argument("--out",
                         help="filename for output file with combined list of files"
                              " for use in metapackage ")
    parser_.add_argument("--build", help="Add build options", action="append", default=[])
    return parser_


def include_list(env, args):
    include = list(env.keys())

    # remove:
    # - all ska3-*-latest
    # - this package from its own requirements
    # - explicit excludes
    include = [pkg for pkg in include if
               not re.match('ska3-\S+-latest', pkg)
               and pkg != args.name
               and pkg not in args.exclude
               ]

    # remove ska3 packages if required
    # (this is to remove ska3-flight-latest from ska3-core)
    if args._in is not None:
        with open(args._in) as fh:
            meta = yaml.load(fh, Loader=yaml.SafeLoader)
            include = [pkg for pkg in include if pkg in meta['requirements']['run']]

    # remove non-ska3 packages if required
    # (this is to remove non-ska3-flight-latest from ska3-flight)
    if args.not_in is not None:
        with open(args.not_in) as fh:
            meta = yaml.load(fh, Loader=yaml.SafeLoader)
            include = [pkg for pkg in include if pkg not in meta['requirements']['run']]

    # add explicit includes
    # (otherwise ska3-flight-latest env minus ska3-flight-latest meta does not include ska3-core)
    for pkg in args.include:
        if pkg in env:
            include.append(pkg)

    include = sorted(include)

    return include


def get_environments(envs, args):
    environments = {}
    logging.info(f'Reading environments for {envs}:')
    for env in envs:
        try:
            platform, filename = env.split('=')
        except ValueError as e:
            logging.info(f' - skipped {env}: ', e)
            raise
        else:
            logging.info(f' + {platform}: {filename}')
            with open(filename) as fh:
                environments[platform] = {p['name']: p for p in json.load(fh)}

            # replace occurrences of ska3-*-latest packages by the current version of ska3-*
            ska3_latest = {}
            for key in environments[platform]:
                if match := re.match('(ska3-\S+)-latest', key):
                    package = match.group(1)
                    logging.info(f'adding {package}')
                    ska3_latest[package] = {
                        'name': package,
                        'version': args.version,
                        'platform': environments[platform][key]['platform']
                    }
            environments[platform].update(ska3_latest)

            environments[platform] = {name: environments[platform][name]
                                      for name in include_list(environments[platform], args)}

    return environments


def main():

    args = parser().parse_args()

    environments = get_environments(args.env, args)
    subtract_environments = get_environments(args.subtract_env, args)

    for platform in environments:
        remove_keys = []
        for package in environments[platform]:
            if (platform in subtract_environments
                    and package in subtract_environments[platform]
                    and subtract_environments[platform][package]['version']
                    == environments[platform][package]['version']):
                remove_keys.append(package)
        for package in remove_keys:
            del environments[platform][package]

    package_names = sorted(set(sum([list(e.keys()) for e in environments.values()], [])))
    all_packages = []
    for p in package_names:
        versions = sorted(set(
            [environments[e][p]['version'].strip() for e in environments if p in environments[e]]
        ))
        for v in versions:
            platforms = sorted([e for e in environments
                                if p in environments[e] and environments[e][p]['version'] == v])
            platforms = [] if len(platforms) == len(environments) else platforms
            all_packages.append({
                'name': p,
                'platforms': ' or '.join(platforms),
                'version': v
            })

    tpl = jinja2.Template(YAML_TPL)
    meta = tpl.render(
        package=args.name,
        version=args.version,
        requirements=all_packages,
        build_options=args.build,
    )
    if args.out:
        out = pathlib.Path(args.out)
        if not out.parent.exists():
            out.parent.mkdir(parents=True)
        with open(out, 'w') as fh:
            fh.write(meta)
    else:
        print(meta)


YAML_TPL = """---
package:
  name: {{ package }}
  version: {{ version }}
{%if build_options %}
build:
  {%- for p in build_options %}
  {{ p }}
  {%- endfor %}
{% endif %}
requirements:
  run:
  {%- for p in requirements %}
    - {{ p.name }} =={{ p.version }}{%if p.platforms %}  # [{{ p.platforms }}]{% endif %}
  {%- endfor %}

"""


if __name__ == "__main__":
    main()