From e15ed3be95f604c11b9522cda94d89675fadca77 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 20 Jan 2024 19:49:41 -0500 Subject: [PATCH] feat(misc): primitive parsing from environment variables --- autotest/test_misc.py | 59 ++++++++++++++++++++++ modflow_devtools/misc.py | 103 ++++++++++++++++++++++++++------------- 2 files changed, 129 insertions(+), 33 deletions(-) diff --git a/autotest/test_misc.py b/autotest/test_misc.py index 7b65641..d51a195 100644 --- a/autotest/test_misc.py +++ b/autotest/test_misc.py @@ -14,6 +14,9 @@ get_packages, has_package, has_pkg, + parse_bool, + parse_float, + parse_int, set_dir, set_env, timed, @@ -280,3 +283,59 @@ def sleep1dec(): cap = capfd.readouterr() print(cap.out) assert re.match(r"sleep1dec took \d+\.\d+ ms", cap.out) + + +def test_parse_bool_default(): + assert parse_bool("TEST_VALUE", default=True) + + +def test_parse_int_default(): + assert parse_int("TEST_VALUE", default=2) == 2 + + +def test_parse_float_default(): + assert parse_float("TEST_VALUE", default=1.1) == 1.1 + + +@pytest.mark.parametrize( + "value, expected", + [ + ("true", True), + ("True", True), + ("TRUE", True), + ("false", False), + ("False", False), + ("FALSE", False), + ("", None), + ("garbage", None), + (1, None), + (1.0, None), + ], +) +def test_parse_bool(value, expected): + with set_env(TEST_VALUE=str(value)): + assert parse_bool("TEST_VALUE") == expected + + +@pytest.mark.parametrize( + "value, expected", + [("1", 1), ("1.", 1), ("1.0", 1), ("1.9", 1), ("", None), ("true", None)], +) +def test_parse_int(value, expected): + with set_env(TEST_VALUE=str(value)): + assert parse_int("TEST_VALUE") == expected + + +@pytest.mark.parametrize( + "value, expected", + [ + ("1", 1.0), + ("", None), + ("1.", 1.0), + ("1.0", 1.0), + ("1.11", 1.11), + ], +) +def test_parse_float(value, expected): + with set_env(TEST_VALUE=str(value)): + assert parse_float("TEST_VALUE") == expected diff --git a/modflow_devtools/misc.py b/modflow_devtools/misc.py index 4915d8c..0ba16bf 100644 --- a/modflow_devtools/misc.py +++ b/modflow_devtools/misc.py @@ -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 @@ -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. @@ -486,3 +454,72 @@ def call(): return res return _timed + + +def parse_bool(name: str, default: Optional[bool] = None) -> Optional[bool]: + """Try to parse a boolean from the environment variable with the given name, + return None or an optional default value if it doesn't exist or parsing fails.""" + try: + v = literal_eval(environ.get(name).lower().title()) + except: + return default + if isinstance(v, bool): + return v + return None + + +def parse_int(name: str, default: Optional[int] = None) -> Optional[int]: + """Try to parse an integer from the environment variable with the given name, + return None or an optional default value if it doesn't exist or parsing fails.""" + try: + v = literal_eval(environ.get(name)) + except: + return default + if isinstance(v, (int, float)): + return int(v) + return None + + +def parse_float(name: str, default: Optional[float] = None) -> Optional[float]: + """Try to parse a float from the environment variable with the given name, + return None or an optional default value if it doesn't exist or parsing fails.""" + try: + v = literal_eval(environ.get(name)) + except: + return default + if isinstance(v, (int, float)): + return float(v) + return None + + +@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]