diff --git a/marimo/_convert/ipynb/to_ir.py b/marimo/_convert/ipynb/to_ir.py index 63e1f262965..1b31401bfee 100644 --- a/marimo/_convert/ipynb/to_ir.py +++ b/marimo/_convert/ipynb/to_ir.py @@ -407,9 +407,18 @@ def magic_env(source: str, command: str) -> str: os.environ['VAR_NAME'] = 'VALUE' """ - del command - _key, value = source.split("=", 1) - return f"import os\nos.environ[{_key!r}] = {value!r}" + if "=" not in source: + return magic_remove(source, command) + key, value = source.split("=", 1) + key = key.strip() + value = value.strip() + if ( + len(value) >= 2 + and value[0] == value[-1] + and value[0] in ("'", '"') + ): + value = value[1:-1] + return f"import os\nos.environ[{key!r}] = {value!r}" def comment_out_code(source: str) -> str: if source.strip(): diff --git a/tests/_convert/ipynb/test_ipynb_to_ir.py b/tests/_convert/ipynb/test_ipynb_to_ir.py index aff574ce76d..d9cbf0880ba 100644 --- a/tests/_convert/ipynb/test_ipynb_to_ir.py +++ b/tests/_convert/ipynb/test_ipynb_to_ir.py @@ -407,6 +407,25 @@ def test_transform_magic_commands_complex(): assert result == expected +def test_transform_magic_env_strips_quotes_and_spacing(): + sources = [ + '%env VAR="hello world"', + "%env PATH", + "%env FOO=bar", + "%env VAR = value", + ] + result = transform_magic_commands(sources) + assert result == [ + "import os\nos.environ['VAR'] = 'hello world'", + ( + "# magic command not supported in marimo; please file an issue to add support\n" + "# %env PATH" + ), + "import os\nos.environ['FOO'] = 'bar'", + "import os\nos.environ['VAR'] = 'value'", + ] + + def test_transform_exclamation_mark_complex(): sources = [ "!pip install package1 package2",