<commit_before>import sys
import os
from functools import partial, lru_cache
from pathlib import Path

import click
import sh
import yaml

from vantage import utils


@click.pass_context
def task_cmd(ctx, path, args):
    env = ctx.obj
    utils.loquacious(f"Running task in {path}")
    meta = load_meta(path)
    env = update_env(meta, env)
    run_required = get_flag(
        env=bool(env.get("VG_RUN_REQUIRED")),
        yml=meta.get("run-required"),
        default=False,
    )
    utils.loquacious(f"  Run required? {'YES' if run_required else 'NO'}")
    if run_required:
        for required in meta.get("requires", []):
            utils.loquacious(f"  Running required task: {required}")
            t = get_task(env, required)
            resp = t.invoke(ctx)
            utils.loquacious(f"  Got resp: {resp}")
    try:
        tty_in = False
        if meta.get("image"):
            utils.loquacious(f"  Spinning up docker image")
            utils.loquacious(f"  Path is: {os.environ.get('PATH')}")
            cmd = sh.Command("docker")
            image = meta.get("image")
            run_args = [
                "run",
                "--volume",
                f"{str(path)}:/vg-task",
                "--label",
                "vantage",
                "--label",
                "vantage-task",
            ]
            if isinstance(image, dict):
                tty_in = bool(image.get("tty", False))
                tag = insert_env_vals(image.pop("tag"), env, args)
                for k, v in image.items():
                    if isinstance(v, list):
                        for w in v:
                            run_args += [
                                f"--{k}",
                                insert_env_vals(w, env, args),
                            ]
                    elif isinstance(v, bool):
                        if v:
                            run_args += [f"--{k}"]
                    else:
                        run_args += [f"--{k}", insert_env_vals(v, env, args)]
            else:
                tag = image
                run_args += ["--rm"]
            for e in env.keys():
                run_args += ["--env", e]
            run_args += [tag, "/vg-task"]
            args = run_args + list(args)
        else:
            utils.loquacious(f"  Passing task over to sh")
            env["PATH"] = os.environ.get("PATH", "")
            cmd = sh.Command(str(path))
        utils.loquacious(f"Running command {cmd} with args {args}")
        cmd(
            *args,
            _fg=True,
            _tty_in=tty_in,
            _out_bufsize=0,
            _env=env,
            _cwd=env["VG_APP_DIR"],
        )
    except sh.ErrorReturnCode as erc:

        utils.loquacious(
            f"  Something went wrong, returned exit code {erc.exit_code}"
        )
        return sys.exit(erc.exit_code)


def insert_env_vals(str, env, args):
    for k, v in env.items():
        needle = f"${k}"
        if needle in str:
            str = str.replace(needle, v)
    for i, v in enumerate(args):
        needle = f"${i}"
        if needle in str:
            str = str.replace(needle, v)
    return str


@lru_cache()
def load_meta(path):
    utils.loquacious(f"  Loading meta from task file")
    with path.open() as f:
        content = f.read()
        sep = None
        for line in content.splitlines():
            if "---" in line:
                sep = line
                break
        if sep is None:
            utils.loquacious(f"  No meta found")
            return {}
        else:
            _, meta, script = content.split(sep, 2)
            comment_marker = sep.replace("---", "")
            utils.loquacious(f"  Meta commented out using '{comment_marker}'")
            meta = meta.replace(f"\n{comment_marker}", "\n")
            return yaml.load(meta, Loader=yaml.SafeLoader)


def update_env(meta, env):
    overrides = meta.get("overrides")
    if overrides is not None:
        if env["VG_VERBOSE"]:
            utils.loquacious("  Updating env with override vars in task meta")
            for key, val in overrides.items():
                utils.loquacious(f"    {key}={val}")
        env.update(overrides)
    defaults = meta.get("defaults")
    if defaults is not None:
        if env["VG_VERBOSE"]:
            utils.loquacious("  Updating env with default vars in task meta")
            for key, val in defaults.items():
                utils.loquacious(f"    {key}={val}")
        defaults.update(env)
        env = defaults
    return env


def get_flag(env, yml, default):
    if env is None:
        if yml is None:
            return default<commit_msg>Stringify yaml vals before doing replce<commit_after>import sys
import os
from functools import partial, lru_cache
from pathlib import Path

import click
import sh
import yaml

from vantage import utils


@click.pass_context
def task_cmd(ctx, path, args):
    env = ctx.obj
    utils.loquacious(f"Running task in {path}")
    meta = load_meta(path)
    env = update_env(meta, env)
    run_required = get_flag(
        env=bool(env.get("VG_RUN_REQUIRED")),
        yml=meta.get("run-required"),
        default=False,
    )
    utils.loquacious(f"  Run required? {'YES' if run_required else 'NO'}")
    if run_required:
        for required in meta.get("requires", []):
            utils.loquacious(f"  Running required task: {required}")
            t = get_task(env, required)
            resp = t.invoke(ctx)
            utils.loquacious(f"  Got resp: {resp}")
    try:
        tty_in = False
        if meta.get("image"):
            utils.loquacious(f"  Spinning up docker image")
            utils.loquacious(f"  Path is: {os.environ.get('PATH')}")
            cmd = sh.Command("docker")
            image = meta.get("image")
            run_args = [
                "run",
                "--volume",
                f"{str(path)}:/vg-task",
                "--label",
                "vantage",
                "--label",
                "vantage-task",
            ]
            if isinstance(image, dict):
                tty_in = bool(image.get("tty", False))
                tag = insert_env_vals(image.pop("tag"), env, args)
                for k, v in image.items():
                    if isinstance(v, list):
                        for w in v:
                            run_args += [
                                f"--{k}",
                                insert_env_vals(w, env, args),
                            ]
                    elif isinstance(v, bool):
                        if v:
                            run_args += [f"--{k}"]
                    else:
                        run_args += [f"--{k}", insert_env_vals(v, env, args)]
            else:
                tag = image
                run_args += ["--rm"]
            for e in env.keys():
                run_args += ["--env", e]
            run_args += [tag, "/vg-task"]
            args = run_args + list(args)
        else:
            utils.loquacious(f"  Passing task over to sh")
            env["PATH"] = os.environ.get("PATH", "")
            cmd = sh.Command(str(path))
        utils.loquacious(f"Running command {cmd} with args {args}")
        cmd(
            *args,
            _fg=True,
            _tty_in=tty_in,
            _out_bufsize=0,
            _env=env,
            _cwd=env["VG_APP_DIR"],
        )
    except sh.ErrorReturnCode as erc:

        utils.loquacious(
            f"  Something went wrong, returned exit code {erc.exit_code}"
        )
        return sys.exit(erc.exit_code)


def insert_env_vals(haystack, env, args):
    for k, v in env.items():
        needle = f"${k}"
        if needle in haystack:
            haystack = haystack.replace(needle, str(v))
    for i, v in enumerate(args):
        needle = f"${i}"
        if needle in haystack:
            haystack = haystack.replace(needle, str(v))
    return haystack


@lru_cache()
def load_meta(path):
    utils.loquacious(f"  Loading meta from task file")
    with path.open() as f:
        content = f.read()
        sep = None
        for line in content.splitlines():
            if "---" in line:
                sep = line
                break
        if sep is None:
            utils.loquacious(f"  No meta found")
            return {}
        else:
            _, meta, script = content.split(sep, 2)
            comment_marker = sep.replace("---", "")
            utils.loquacious(f"  Meta commented out using '{comment_marker}'")
            meta = meta.replace(f"\n{comment_marker}", "\n")
            return yaml.load(meta, Loader=yaml.SafeLoader)


def update_env(meta, env):
    if meta.get("env"):
        env_file = meta["env"]
        utils.loquacious("  Using env {env_file} instead")
        return load_env(env_file, current=env)
    overrides = meta.get("overrides")
    if overrides is not None:
        if env["VG_VERBOSE"]:
            utils.loquacious("  Updating env with override vars in task meta")
            for key, val in overrides.items():
                utils.loquacious(f"    {key}={val}")
        env.update(overrides)
    defaults = meta.get("defaults")
    if defaults is not None:
        if env["VG_VERBOSE"]:
            utils.loquacious("  Updating env with default vars in task meta")
            for key, val in defaults.items():
                utils.loquacious(f"    {key}={val}")
        defaults.update(env)
        env = defaults
    return env


def load_env(name_or_path, current):
    new_env = None
    path = Path(name_or_path)
    if path.is_file():
        new_env = utils.load_env_from_file(path)
    else:
        env_dir = current.get("VG_ENV_DIR")
        if env_dir:
            path = Path(env_dir) / name_or_path
            if path.is_file():
                new_env = utils.load_env_from_file(path)
    if new_env is None:
        raise click.ClickException(
            f"The env file '{name_or_path}' does not exist"
        )
    for k, v in current.items():
        if k.startswith("VG_") and k not in new_env:
            new_env[k] = v
    return new_env


def get_flag(env, yml, default):
    if env is None:
        if yml is None:
            return default
{'fraction_non_alphanumeric': 0.11897516442524786, 'fraction_numerical': 0.0003926573083341514, 'mean_word_length': 1.5572289156626506, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5063261', 'n_tokens_mistral': 3091, 'n_tokens_neox': 3023, 'n_words': 782}