Add expandvars and respective tests. Added lit to pixi.toml, and `m…

…agic 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 <[email protected]> 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 <[email protected]>
modularml · Nov 4, 2024 · 884cb52 · 884cb52
1 parent 388eaa4
commit 884cb52
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 32 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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 → stdlib/test/os/test_env.mojo b/stdlib/test/os/test_getenv_setenv.mojo → 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()