Assignment to env var in Jupyter Notebook doesn't round-trip (#2642)

closes #2641

diff --git a/CHANGES.md b/CHANGES.md
index 94d1c69259a413e2af37ee1d453c4a16f19b69ff..d81c1fd4fc4c4ddcdc24a63f6f4d58122604743b 100644 (file)
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -6,6 +6,7 @@
 
 - Fixed Python 3.10 support on platforms without ProcessPoolExecutor (#2631)
 - Fixed `match` statements with open sequence subjects, like `match a, b:` (#2639)
+- Fixed assignment to environment variables in Jupyter Notebooks (#2642)
 
 ## 21.11b1
 

diff --git a/src/black/handle_ipynb_magics.py b/src/black/handle_ipynb_magics.py
index 2fe6739209dd626597a43011b8bc576ae203cc9a..5807dac14d0286cb77fb36dd8af0820f3bc2ab75 100644 (file)
--- a/src/black/handle_ipynb_magics.py
+++ b/src/black/handle_ipynb_magics.py
@@ -403,20 +403,28 @@ class MagicFinder(ast.NodeVisitor):
         For example,
 
             black_version = !black --version
+            env = %env var
 
-        would have been transformed to
+        would have been (respectively) transformed to
 
             black_version = get_ipython().getoutput('black --version')
+            env = get_ipython().run_line_magic('env', 'var')
 
-        and we look for instances of the latter.
+        and we look for instances of any of the latter.
         """
-        if (
-            isinstance(node.value, ast.Call)
-            and _is_ipython_magic(node.value.func)
-            and node.value.func.attr == "getoutput"
-        ):
-            (arg,) = _get_str_args(node.value.args)
-            src = f"!{arg}"
+        if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
+            args = _get_str_args(node.value.args)
+            if node.value.func.attr == "getoutput":
+                src = f"!{args[0]}"
+            elif node.value.func.attr == "run_line_magic":
+                src = f"%{args[0]}"
+                if args[1]:
+                    src += f" {args[1]}"
+            else:
+                raise AssertionError(
+                    "Unexpected IPython magic {node.value.func.attr!r} found. "
+                    "Please report a bug on https://github.com/psf/black/issues."
+                ) from None
             self.magics[node.value.lineno].append(
                 OffsetAndMagic(node.value.col_offset, src)
             )
@@ -451,7 +459,6 @@ class MagicFinder(ast.NodeVisitor):
                 else:
                     src = f"%{args[0]}"
                     if args[1]:
-                        assert src is not None
                         src += f" {args[1]}"
             elif node.value.func.attr == "system":
                 src = f"!{args[0]}"

diff --git a/tests/test_ipynb.py b/tests/test_ipynb.py
index ba460074e9a8650967a8f261df881a5d8492da2b..5ecdf00b7bce8d50637fff3364ff9637ca84d206 100644 (file)
--- a/tests/test_ipynb.py
+++ b/tests/test_ipynb.py
@@ -90,6 +90,10 @@ def test_cell_magic_noop() -> None:
             id="Line magic with argument",
         ),
         pytest.param("%time\n'foo'", '%time\n"foo"', id="Line magic without argument"),
+        pytest.param(
+            "env =  %env var", "env = %env var", id="Assignment to environment variable"
+        ),
+        pytest.param("env =  %env", "env = %env", id="Assignment to magic"),
     ),
 )
 def test_magic(src: str, expected: str) -> None: