veros_copy_setup.py

#!/usr/bin/env python

import os
import shutil
import datetime
import functools
import textwrap
import importlib

import click
import entrypoints


SETUPDIR_ENVVAR = "VEROS_SETUP_DIR"
IGNORE_PATTERNS = ["__init__.py", "*.pyc", "__pycache__"]
SETUPS = {}

setup_dirs = []


for e in entrypoints.get_group_all("veros.setup_dirs"):
    # - hack to prevent actually importing plugins -
    # we can only find the location of top-level modules
    # so assume that foo.bar.baz always resolves to foo/bar/baz
    base_module, *submodules = e.module_name.split(".")

    modpath = os.path.join(os.path.dirname(importlib.util.find_spec(base_module).origin), *submodules)
    setup_dirs.append(modpath)

for setup_dir in os.environ.get(SETUPDIR_ENVVAR, "").split(";"):
    if os.path.isdir(setup_dir):
        setup_dirs.append(setup_dir)

# populate {setup_name: path} mapping
for setup_dir in setup_dirs:
    for setup in os.listdir(setup_dir):
        setup_path = os.path.join(setup_dir, setup)

        if not os.path.isdir(setup_path):
            continue

        if setup.startswith(("_", ".")):
            continue

        SETUPS[setup] = setup_path

SETUP_NAMES = sorted(SETUPS.keys())


def rewrite_main_file(target_file, setup_name):
    from veros import __version__ as veros_version

    current_date = datetime.datetime.utcnow()
    header_str = textwrap.dedent(
        f'''
        """
        This Veros setup file was generated by

           $ veros copy-setup {setup_name}

        on {current_date:%Y-%m-%d %H:%M:%S} UTC.
        """

        __VEROS_VERSION__ = {veros_version!r}

        if __name__ == "__main__":
            raise RuntimeError(
                "Veros setups cannot be executed directly. "
                f"Try `veros run {{__file__}}` instead."
            )

        # -- end of auto-generated header, original file below --
    '''
    ).strip()

    with open(target_file, "r") as f:
        orig_contents = f.readlines()

    shebang = None
    if orig_contents[0].startswith("#!"):
        shebang = orig_contents[0]
        orig_contents = orig_contents[1:]

    with open(target_file, "w") as f:
        if shebang is not None:
            f.write(shebang + "\n")

        f.write(header_str + "\n\n")
        f.writelines(orig_contents)


def copy_setup(setup, to=None):
    """Copy a standard setup to another directory.

    Available setups:

        {setups}

    Example:

        $ veros copy-setup global_4deg --to ~/veros-setups/4deg-lowfric

    Further directories containing setup templates can be added to this command
    via the {setup_envvar} environment variable.
    """
    if to is None:
        to = os.path.join(os.getcwd(), setup)

    if os.path.exists(to):
        raise RuntimeError("Target directory must not exist")

    to_parent = os.path.dirname(os.path.realpath(to))

    if not os.path.exists(to_parent):
        os.makedirs(to_parent)

    ignore = shutil.ignore_patterns(*IGNORE_PATTERNS)
    shutil.copytree(SETUPS[setup], to, ignore=ignore)

    main_setup_file = os.path.join(to, f"{setup}.py")
    rewrite_main_file(main_setup_file, setup)


copy_setup.__doc__ = copy_setup.__doc__.format(setups=", ".join(SETUP_NAMES), setup_envvar=SETUPDIR_ENVVAR)


@click.command("veros-copy-setup")
@click.argument("setup", type=click.Choice(SETUP_NAMES), metavar="SETUP")
@click.option(
    "--to",
    required=False,
    default=None,
    type=click.Path(dir_okay=False, file_okay=False, writable=True),
    help=("Target directory, must not exist " "(default: copy to current working directory)"),
)
@functools.wraps(copy_setup)
def cli(*args, **kwargs):
    copy_setup(*args, **kwargs)