Source code for aviary.interface.download_models

import argparse
import os
import shutil
from pathlib import Path

from aviary.utils.functions import get_model




def save_file(aviary_path: Path, outdir: Path, verbose=False) -> Path:
    """Saves the file or folder specified into the output directory, creating directories as needed."""
    outdir.mkdir(parents=True, exist_ok=True)
    if aviary_path.is_dir():
        if verbose:
            print(aviary_path, 'is a directory, getting all files')
        outdir = outdir.joinpath(aviary_path.stem)
        outdir.mkdir(exist_ok=True)
        for file in next(os.walk(aviary_path))[-1]:
            if verbose:
                print('copying', str(aviary_path / file), 'to', str(outdir / file))
            shutil.copy2(aviary_path / file, outdir)
    else:
        if verbose:
            print('copying', str(aviary_path), 'to', str(outdir / aviary_path.name))
        shutil.copy2(aviary_path, outdir)
    return outdir



def _setup_hangar_parser(parser: argparse.ArgumentParser):
    def_outdir = os.path.join(os.getcwd(), 'aviary_models')
    parser.add_argument(
        'input_decks',
        metavar='indecks',
        type=str,
        nargs='+',
        help='Name of file or folder to download from Aviary/models',
    )
    parser.add_argument(
        '-o',
        '--outdir',
        default=def_outdir,
        help='Directory to write outputs. Defaults to aviary_models in the current directory.',
    )
    parser.add_argument(
        '-v',
        '--verbose',
        action='store_true',
        help='Enable verbose outputs',
    )


def _exec_hangar(args, user_args):
    input_decks = []
    for input_deck in args.input_decks:
        input_decks.append(get_model(input_deck))

    for input_deck in input_decks:
        save_file(input_deck, Path(args.outdir), args.verbose)