Source code for chariot.context

import os
import shutil

from chariot.config import getLogger, reload, settings
from chariot_api import chariot_dir

log = getLogger(__name__)

env_file = os.path.expanduser(settings.model_config["env_file"])
secret_file = os.path.expanduser(os.path.join(chariot_dir, "client_secret"))
context_path = os.path.expanduser(os.path.join(chariot_dir, ".current_context"))




def get_contexts(detailed=False):
    if not detailed:
        return [f[5:] for f in os.listdir(chariot_dir) if f.startswith(".env-")]
    envs = []
    for f in os.listdir(chariot_dir):
        if not f.startswith(".env-"):
            continue
        env = {
            k: v
            for k, _, v in (
                line.strip().partition("=") for line in open(f"{chariot_dir}/{f}").readlines()
            )
        }
        envs.append({"name": f[5:], **env})
    return envs





def get_conf_files_from_context(context):
    context_env_file = f"{env_file}-{context}"
    context_secret_file = f"{secret_file}-{context}"
    return context_env_file, context_secret_file





def use_context(context):
    context_env_file, context_secret_file = get_conf_files_from_context(context)
    shutil.copy(context_env_file, env_file)
    shutil.copy(context_secret_file, secret_file)

    with open(context_path, "w") as f:
        f.write(context)

    reload()





def clear():
    os.remove(env_file)
    os.remove(secret_file)
    os.remove(context_path)
    reload()





def current_context():
    try:
        with open(context_path) as f:
            return f.read()
    except FileNotFoundError:
        log.warning(
            "No current context found, please create-context and use-context to ride the Chariot"
        )
        return None





def create_context(
    context: str,
    *,
    client_id: str = None,
    secret: str = None,
    url: str = None,
    user_id: str = None,
    scope: str = "user",
    tls_verify: bool = True,
    overwrite: bool = True,
    **kwargs,
):
    context_env_file, context_secret_file = get_conf_files_from_context(context)
    if os.path.exists(context_env_file) and not overwrite:
        log.error(f"context '{context}' already exists and overwrite=False, will not update")
        return

    with open(context_env_file, "w") as f:
        f.write(f"USER_ID={user_id}\n")
        f.write(f"CLIENT_ID={client_id}\n")
        f.write(f"CLIENT_SCOPE={scope}\n")
        f.write(f"BASE_URL={url or settings.base_url}\n")
        f.write(f"VERIFY_SSL={tls_verify}\n")
        f.write(f"EXTRA_ARGS={kwargs}\n")

    with open(context_secret_file, "w") as f:
        f.write(secret)





def delete_context(context):
    context_env_file, context_secret_file = get_conf_files_from_context(context)
    os.remove(context_env_file)
    os.remove(context_secret_file)