""" Environment-related utilities """

import os

__author__ = "Vince Reuter"
__email__ = "vreuter@virginia.edu"

__all__ = ["TmpEnv"]


class TmpEnv(object):
    """Temporary environment variable setting."""

    def __init__(self, overwrite=False, **kwargs):
        if not overwrite:
            already_set = [k for k, v in kwargs.items() if os.getenv(k, v) != v]
            if already_set:
                msg = "{} variable(s) already set: {}".format(
                    len(already_set), ", ".join(already_set)
                )
                raise ValueError(msg)
        self._kvs = kwargs

    def __enter__(self):
        for k, v in self._kvs.items():
            os.environ[k] = v
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        for k in self._kvs:
            try:
                del os.environ[k]
            except KeyError:
                pass