From 884cb525ea084db62be5cee91da1e1e50c08e850 Mon Sep 17 00:00:00 2001 From: Mikhail Tavarez Date: Sun, 3 Nov 2024 13:00:07 -0600 Subject: [PATCH] Add `expandvars` and respective tests. Added lit to pixi.toml, and `magic run` to pre-commit configuration. squash code a bit by using write add tests Cleanup with formatting and sign add missing newline Signed-off-by: Mikhail Tavarez updated changelog updated changelog Add unsetenv and use a context manager to unset test variables. Add renamed test_env.mojo Fix unsetenv result in EnvVar context manager. Signed-off-by: Mikhail Tavarez --- stdlib/src/os/__init__.mojo | 2 +- stdlib/src/os/env.mojo | 20 ++++++ stdlib/src/os/path/path.mojo | 10 ++- stdlib/test/os/path/test_expandvars.mojo | 65 ++++++++++++------- ...{test_getenv_setenv.mojo => test_env.mojo} | 10 ++- 5 files changed, 75 insertions(+), 32 deletions(-) rename stdlib/test/os/{test_getenv_setenv.mojo => test_env.mojo} (84%) diff --git a/stdlib/src/os/__init__.mojo b/stdlib/src/os/__init__.mojo index fe813bb7d9b..b55400cad34 100644 --- a/stdlib/src/os/__init__.mojo +++ b/stdlib/src/os/__init__.mojo @@ -13,7 +13,7 @@ """Implements the os package.""" from .atomic import Atomic -from .env import getenv, setenv +from .env import getenv, setenv, unsetenv from .fstat import lstat, stat, stat_result from .os import ( SEEK_CUR, diff --git a/stdlib/src/os/env.mojo b/stdlib/src/os/env.mojo index 6feb36e79b4..0500e81bbde 100644 --- a/stdlib/src/os/env.mojo +++ b/stdlib/src/os/env.mojo @@ -51,6 +51,26 @@ fn setenv(name: String, value: String, overwrite: Bool = True) -> Bool: return status == 0 +fn unsetenv(name: String) -> Bool: + """Unsets an environment variable. + + Constraints: + The function only works on macOS or Linux and returns False otherwise. + + Args: + name: The name of the environment variable. + + Returns: + True if unsetting the variable succeeded. Otherwise, False is returned. + """ + alias os_is_supported = os_is_linux() or os_is_macos() + if not os_is_supported: + return False + + var status = external_call["unsetenv", Int32](name.unsafe_ptr()) + return status == 0 + + fn getenv(name: String, default: String = "") -> String: """Returns the value of the given environment variable. diff --git a/stdlib/src/os/path/path.mojo b/stdlib/src/os/path/path.mojo index fcd502433cb..ffdf5e0bdc2 100644 --- a/stdlib/src/os/path/path.mojo +++ b/stdlib/src/os/path/path.mojo @@ -426,9 +426,7 @@ fn _is_shell_special_variable(byte: Byte) -> Bool: ord("8"), ord("9"), ) - if int(byte) in shell_variables: - return True - return False + return int(byte) in shell_variables fn _is_alphanumeric(byte: Byte) -> Bool: @@ -492,16 +490,16 @@ fn _parse_variable_name(bytes: Span[Byte]) -> Tuple[String, Int]: fn expandvars[PathLike: os.PathLike, //](path: PathLike) -> String: """Replaces `${var}` or `$var` in the path with values from the current environment variables. - Undefined variables should be left alone. + Malformed variable names and references to non-existing variables are left unchanged. Parameters: PathLike: The type conforming to the os.PathLike trait. Args: - path: The path to expand. + path: The path that is being expanded. Returns: - The input path with environment variables expanded. + The expanded path. """ var path_str = path.__fspath__() var bytes = path_str.as_bytes() diff --git a/stdlib/test/os/path/test_expandvars.mojo b/stdlib/test/os/path/test_expandvars.mojo index e661675b6aa..fa11e4f9f02 100644 --- a/stdlib/test/os/path/test_expandvars.mojo +++ b/stdlib/test/os/path/test_expandvars.mojo @@ -17,36 +17,53 @@ from os.path import expandvars from testing import assert_equal +@value +struct EnvVar: + var name: String + + fn __init__(inout self, name: String, value: String) -> None: + self.name = name + _ = os.setenv(name, value) + + fn __enter__(self) -> Self: + return self + + fn __exit__(self) -> None: + _ = os.unsetenv(self.name) + + def test_expansion(): - _ = os.setenv("TEST_VAR", "World") - assert_equal(expandvars("Hello $TEST_VAR!"), "Hello World!") - assert_equal(expandvars("漢字 $TEST_VAR🔥!"), "漢字 World🔥!") - assert_equal(expandvars("$TEST_VAR/path/to/file"), "World/path/to/file") - - _ = os.setenv("UNICODE_TEST_VAR", "漢字🔥") - assert_equal(expandvars("Hello $UNICODE_TEST_VAR!"), "Hello 漢字🔥!") - assert_equal(expandvars("漢字 $UNICODE_TEST_VAR🔥!"), "漢字 漢字🔥🔥!") - assert_equal( - expandvars("$UNICODE_TEST_VAR/path/to/file"), "漢字🔥/path/to/file" - ) + with EnvVar("TEST_VAR", "World"): + assert_equal(expandvars("Hello $TEST_VAR!"), "Hello World!") + assert_equal(expandvars("漢字 $TEST_VAR🔥!"), "漢字 World🔥!") + assert_equal(expandvars("$TEST_VAR/path/to/file"), "World/path/to/file") + + with EnvVar("UNICODE_TEST_VAR", "漢字🔥"): + assert_equal(expandvars("Hello $UNICODE_TEST_VAR!"), "Hello 漢字🔥!") + assert_equal(expandvars("漢字 $UNICODE_TEST_VAR🔥!"), "漢字 漢字🔥🔥!") + assert_equal( + expandvars("$UNICODE_TEST_VAR/path/to/file"), "漢字🔥/path/to/file" + ) def test_braced_expansion(): - _ = os.setenv("BRACE_VAR", "World") - assert_equal(expandvars("Hello ${BRACE_VAR}!"), "Hello World!") - assert_equal(expandvars("漢字 ${BRACE_VAR}🔥!"), "漢字 World🔥!") - assert_equal(expandvars("${BRACE_VAR}/path/to/file"), "World/path/to/file") - - _ = os.setenv("UNICODE_BRACE_VAR", "漢字🔥") - assert_equal(expandvars("Hello ${UNICODE_BRACE_VAR}!"), "Hello 漢字🔥!") - assert_equal(expandvars("漢字 ${UNICODE_BRACE_VAR}🔥!"), "漢字 漢字🔥🔥!") - assert_equal( - expandvars("${UNICODE_BRACE_VAR}/path/to/file"), "漢字🔥/path/to/file" - ) + with EnvVar("BRACE_VAR", "World"): + assert_equal(expandvars("Hello ${BRACE_VAR}!"), "Hello World!") + assert_equal(expandvars("漢字 ${BRACE_VAR}🔥!"), "漢字 World🔥!") + assert_equal( + expandvars("${BRACE_VAR}/path/to/file"), "World/path/to/file" + ) + + with EnvVar("UNICODE_BRACE_VAR", "漢字🔥"): + assert_equal(expandvars("Hello ${UNICODE_BRACE_VAR}!"), "Hello 漢字🔥!") + assert_equal(expandvars("漢字 ${UNICODE_BRACE_VAR}🔥!"), "漢字 漢字🔥🔥!") + assert_equal( + expandvars("${UNICODE_BRACE_VAR}/path/to/file"), "漢字🔥/path/to/file" + ) def test_unset_expansion(): - # Unset variables should be expanded to an empty string. + # Unset variables should not be expanded. assert_equal( expandvars("Hello $NONEXISTENT_VAR!"), "Hello $NONEXISTENT_VAR!" ) @@ -56,7 +73,7 @@ def test_unset_expansion(): def test_dollar_sign(): - # A lone dollar sign should not be expanded. + # A lone `$` should not be expanded. assert_equal(expandvars("A lone $ sign"), "A lone $ sign") # Special shell variables should not be expanded. diff --git a/stdlib/test/os/test_getenv_setenv.mojo b/stdlib/test/os/test_env.mojo similarity index 84% rename from stdlib/test/os/test_getenv_setenv.mojo rename to stdlib/test/os/test_env.mojo index 5c1ffad8328..0d4bd0b6ec0 100644 --- a/stdlib/test/os/test_getenv_setenv.mojo +++ b/stdlib/test/os/test_env.mojo @@ -13,7 +13,7 @@ # REQUIRES: system-linux || system-darwin # RUN: TEST_MYVAR=MyValue %mojo %s -from os import getenv, setenv +from os import getenv, setenv, unsetenv from testing import assert_equal @@ -40,6 +40,14 @@ def test_setenv(): assert_equal(setenv("=", "INVALID", True), False) +def test_unsetenv(): + assert_equal(setenv("NEW_VAR", "FOO", True), True) + assert_equal(getenv("NEW_VAR"), "FOO") + assert_equal(unsetenv("NEW_VAR"), True) + assert_equal(getenv("NEW_VAR"), "") + + def main(): test_getenv() test_setenv() + test_unsetenv()