Skip to content

Commit

Permalink
feat(misc): parse literals from environment variables (#135)
Browse files Browse the repository at this point in the history
  • Loading branch information
wpbonelli authored Jan 21, 2024
1 parent fd21500 commit a9b8019
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 33 deletions.
23 changes: 23 additions & 0 deletions autotest/test_misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import pytest

from modflow_devtools.misc import (
get_env,
get_model_paths,
get_namefile_paths,
get_packages,
Expand Down Expand Up @@ -280,3 +281,25 @@ def sleep1dec():
cap = capfd.readouterr()
print(cap.out)
assert re.match(r"sleep1dec took \d+\.\d+ ms", cap.out)


def test_get_env():
assert get_env("NO_VALUE") is None

with set_env(TEST_VALUE=str(True)):
assert get_env("NO_VALUE", True) == True
assert get_env("TEST_VALUE") == True
assert get_env("TEST_VALUE", default=False) == True
assert get_env("TEST_VALUE", default=1) == 1

with set_env(TEST_VALUE=str(1)):
assert get_env("NO_VALUE", 1) == 1
assert get_env("TEST_VALUE") == 1
assert get_env("TEST_VALUE", default=2) == 1
assert get_env("TEST_VALUE", default=2.1) == 2.1

with set_env(TEST_VALUE=str(1.1)):
assert get_env("NO_VALUE", 1.1) == 1.1
assert get_env("TEST_VALUE") == 1.1
assert get_env("TEST_VALUE", default=2.1) == 1.1
assert get_env("TEST_VALUE", default=False) == False
99 changes: 66 additions & 33 deletions modflow_devtools/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import socket
import sys
import traceback
from ast import literal_eval
from contextlib import contextmanager
from functools import wraps
from importlib import metadata
Expand Down Expand Up @@ -31,39 +32,6 @@ def set_dir(path: PathLike):
print(f"Returned to previous directory: {origin}")


@contextmanager
def set_env(*remove, **update):
"""
Temporarily updates the ``os.environ`` dictionary in-place.
Referenced from https://stackoverflow.com/a/34333710/6514033.
The ``os.environ`` dictionary is updated in-place so that the modification
is sure to work in all situations.
:param remove: Environment variables to remove.
:param update: Dictionary of environment variables and values to add/update.
"""
env = environ
update = update or {}
remove = remove or []

# List of environment variables being updated or removed.
stomped = (set(update.keys()) | set(remove)) & set(env.keys())
# Environment variables and values to restore on exit.
update_after = {k: env[k] for k in stomped}
# Environment variables and values to remove on exit.
remove_after = frozenset(k for k in update if k not in env)

try:
env.update(update)
[env.pop(k, None) for k in remove]
yield
finally:
env.update(update_after)
[env.pop(k) for k in remove_after]


class add_sys_path:
"""
Context manager to add temporarily to the system path.
Expand Down Expand Up @@ -486,3 +454,68 @@ def call():
return res

return _timed


def get_env(name: str, default: object = None) -> Optional[object]:
"""
Try to parse the given environment variable as the type of the given
default value, if one is provided, otherwise any type is acceptable.
If the types of the parsed value and default value don't match, the
default value is returned. The environment variable is parsed as a
Python literal with `ast.literal_eval()`.
Parameters
----------
name : str
The environment variable name
default : object
The default value if the environment variable does not exist
Returns
-------
The value of the environment variable, parsed as a Python literal,
otherwise the default value if the environment variable is not set.
"""
try:
v = environ.get(name)
if isinstance(default, bool):
v = v.lower().title()
v = literal_eval(v)
except:
return default
if default is None:
return v
return v if isinstance(v, type(default)) else default


@contextmanager
def set_env(*remove, **update):
"""
Temporarily updates the ``os.environ`` dictionary in-place.
Referenced from https://stackoverflow.com/a/34333710/6514033.
The ``os.environ`` dictionary is updated in-place so that the modification
is sure to work in all situations.
:param remove: Environment variables to remove.
:param update: Dictionary of environment variables and values to add/update.
"""
env = environ
update = update or {}
remove = remove or []

# List of environment variables being updated or removed.
stomped = (set(update.keys()) | set(remove)) & set(env.keys())
# Environment variables and values to restore on exit.
update_after = {k: env[k] for k in stomped}
# Environment variables and values to remove on exit.
remove_after = frozenset(k for k in update if k not in env)

try:
env.update(update)
[env.pop(k, None) for k in remove]
yield
finally:
env.update(update_after)
[env.pop(k) for k in remove_after]

0 comments on commit a9b8019

Please sign in to comment.