pyparepackage.py

import re
import sys
import yaml
import argparse
import subprocess
from pathlib import Path
from collections import defaultdict
from queue import Queue
from typing import Optional, Generic, Type, TypeVar, MutableMapping, List

Node = TypeVar('Node')
architectures: MutableMapping = {}
command_options: List[str] = ["missing", "buildout"]

class Package(Generic[Node]):
    name: str
    version: str
    full_n_v: str
    prev: Optional[Node]
    variants: int
    install_flags: str
    compiler: str

class Packages(Package):
    def __init__(self) -> None:
        self.length: int = 0
        self.head: Optional[Package] = None

    def push(self, n: str, v: str, variants: int, flag: str, compiler: str):
        pkg_node: Package = Package(
            name=n,
            version=v,
            full_n_v=f'{n}@{v}',
            prev=None,
            variants=variants,
            install_flags=flag,
            compiler=compiler
        )

        self.length += 1
        if not self.head:
            self.head = pkg_node
            return

        pkg_node.prev = self.head
        self.head = pkg_node

    def pop(self):
        self.length = max([0, self.length - 1])
        if self.length == 0:
            head = self.head
            if head:
                self.head = None
                return head.full_n_v, head.variants, head.install_flags, head.compiler
            else:
                return None

        head = self.head
        self.head = head.prev

        return head.full_n_v, head.variants, head.install_flags, head.compiler

class Environment:
    def __init__(self):
        self.requirements: Type[Packages] = None
        self.missing: List[str] = []
        self.compiler: str # to replace with yaml_data

    def generate_requirements(self, compiler: str, dir_main_requested: str):
        self.compiler = compiler # to replace with yaml_data
        packages = Packages()
        
        for _filename in Path(dir_main_requested).iterdir():
            if _filename.suffix == '.yaml':
                _pkg_yaml = _filename.open().read()
                yaml_data = yaml.safe_load(_pkg_yaml)

                for version in yaml_data['versions'].keys():
                    [
                        packages.push(
                            n=_filename.stem,
                            v=version,
                            variants=len(yaml_data['versions'][version]),
                            flag=yaml_data['versions'][version][count],
                            compiler=compiler # to replace with yaml_data
                        )
                        for count in range(len(yaml_data['versions'][version]))
                    ]
        self.requirements = packages

    def find_missing_packages(self, arch: str):
        packages = self.requirements

        spack_command_flags = "-vf --show-full-compiler"
        spack_output = subprocess.run(
            f'~/spack/bin/spack find {spack_command_flags} {self.compiler} arch={arch}',
            capture_output = True,
            text = True,
            shell = True,
            timeout=180,
        )

        spack_found_pkgs = [
            re.split(r'(?=%)[^" "]*', pkg)
            for pkg in list(spack_output.stdout.splitlines())[1:-1]
        ]

        # multiple variants, same version
        installed_pkgs = defaultdict(list)
        for pkg_name_ver, pkg_flags in spack_found_pkgs:
            installed_pkgs[pkg_name_ver].append(pkg_flags)

        missing_packages = Queue(
            maxsize=packages.length
        )
        while packages.length:
            package, variants, flags, compiler = packages.pop()
            if package in installed_pkgs.keys() and variants > 1:
                # test for default no variant
                if all([ flag != '' for flag in installed_pkgs[package]]) and flags == '':
                    missing_packages.put(package_data(package, compiler, flags))
                    continue
                # test variant with no occurence
                if flags != '':
                    present = 0
                    for flag_set in range(0, len(installed_pkgs[package])):
                        if all(flag in installed_pkgs[package][flag_set] for flag in flags.split()):
                            present += 1
                    if not present:
                        missing_packages.put(package_data(package, compiler, flags))
                        continue
            if package not in installed_pkgs.keys():
                missing_packages.put(package_data(package, compiler, flags))
        self.missing = missing_packages

    def show_missing_packages(self):
        show = []
        missing_packages = self.missing

        while not missing_packages.empty():
            pkg_name_versions, pkg_flags = missing_packages.get()
            n = pkg_name_versions[0]
            v = pkg_name_versions[1]
            show.append(f'{n}@{v}: {pkg_flags}')
        return show

    def show_misssing_pkg_build_results(self, arch: str):
        missing_packages = self.missing
        show = []
        matched_arch = [
            _file for _file
            in Path("data/results").iterdir()
            if arch in _file.stem
        ]

        while not missing_packages.empty():
            build_time, build_result_filename = '', ''
            pkg_name_version, pkg_flags = missing_packages.get()
            for _filename in matched_arch:
                if _filename.stem[:len(pkg_name_version[0])] != pkg_name_version[0]:
                    continue
                _pkg_yaml = _filename.open().read()
                yaml_data = yaml.safe_load(_pkg_yaml)
                comp_name = pkg_name_version[0] == yaml_data['package']
                comp_verison = pkg_name_version[1] == yaml_data['version']
                comp_flags = pkg_flags == yaml_data['flags']
                if all([comp_name, comp_verison, comp_flags]):
                    if build_time > yaml_data['end_time']:
                        continue
                    build_time = yaml_data['end_time']
                    build_result_filename = yaml_data['build_name']
            show.append(build_result_filename)
        return show

def package_data(pkg_name_version: str, compiler: str, flags: str):
    return [pkg_name_version.split('@'), " ".join([compiler, flags])]

def check_fileobject(path: str, fileobject: str):
    try:
        if fileobject == 'file' and not Path(path).is_file():
            raise FileNotFoundError(path, fileobject)
        if fileobject == 'directory' and not Path(path).is_dir():
            raise FileNotFoundError(path, fileobject)
    except FileNotFoundError as err:
        p, o = err.args
        print(f'Invalid {o}, missing {p} in:')
        print(Path.cwd())
        exit(1)

def main(args):
    search_cmd = args.search_command
    base_dir = args.main_dir
    build_compiler = f"%{args.compiler}"
    build_archs = args.archs
    check_fileobject(base_dir, 'directory')

    for arch in build_archs:
        architectures[arch] = Environment()
        architectures[arch].generate_requirements(build_compiler, base_dir)
        architectures[arch].find_missing_packages(arch)

    if search_cmd == "missing":
        for arch in build_archs:
            print(arch)
            for pkg in architectures[arch].show_missing_packages():
                print(pkg)
            print("\n")
    if search_cmd == "buildout":
        for arch in build_archs:
            print(arch)
            for result in architectures[arch].show_misssing_pkg_build_results(arch=arch):
                print(result)
            print("\n")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('search_command')
    parser.add_argument('main_dir')
    parser.add_argument('compiler', default="gcc")
    parser.add_argument('--archs', nargs="*")
    args = parser.parse_args()
    exit(main(args=args))