From e0253080b0d2b61bf2105a2f5afdf5173e33d0e5 Mon Sep 17 00:00:00 2001
From: Marco Edward Gorelli <marcogorelli@protonmail.com>
Date: Fri, 26 Nov 2021 16:14:57 +0000
Subject: [PATCH]  Assignment to env var in Jupyter Notebook doesn't round-trip
 (#2642)

closes #2641
---
 CHANGES.md                       |  1 +
 src/black/handle_ipynb_magics.py | 27 +++++++++++++++++----------
 tests/test_ipynb.py              |  4 ++++
 3 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 94d1c69..d81c1fd 100644
--- 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 2fe6739..5807dac 100644
--- 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 ba46007..5ecdf00 100644
--- 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:
-- 
2.39.5