Repository: psf/black

Issue:
Assignment to env var in Jupyter Notebook doesn't round-trip
<!--
Please make sure that the bug is not already fixed either in newer versions or the
current development version. To confirm this, you have three options:

1. Update Black's version if a newer release exists: `pip install -U black`
2. Use the online formatter at <https://black.vercel.app/?version=main>, which will use
   the latest main branch.
3. Or run _Black_ on your machine:
   - create a new virtualenv (make sure it's the same Python version);
   - clone this repository;
   - run `pip install -e .[d,python2]`;
   - run `pip install -r test_requirements.txt`
   - make sure it's sane by running `python -m pytest`; and
   - run `black` like you did last time.
-->

**Describe the bug**

<!-- A clear and concise description of what the bug is. -->

A cell with
```
env =  %env var
```
, run `black` on the notebook. Originally reported here https://github.com/nbQA-dev/nbQA/issues/679 (thanks @danielsparing !)

**To Reproduce**

<!--
Minimal steps to reproduce the behavior with source code and Black's configuration.
-->


The cell gets transformed to
```
var = get_ipython().run_line_magic("env", "var")
```

**Expected behavior**

```
env = %env var
```
**Environment**

<!-- Please complete the following information: -->

- Black's version: <!-- e.g. [main] -->
- OS and Python version: <!-- e.g. [Linux/Python 3.7.4rc1] -->

**Additional context**

<!-- Add any other context about the problem here. -->

Will open a PR to address this now


Discussion:

Resolving patch:
diff --git a/src/black/handle_ipynb_magics.py b/src/black/handle_ipynb_magics.py
--- a/src/black/handle_ipynb_magics.py
+++ b/src/black/handle_ipynb_magics.py
@@ -403,20 +403,28 @@ def visit_Assign(self, node: ast.Assign) -> None:
         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 @@ def visit_Expr(self, node: ast.Expr) -> None:
                 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]}"


Test patch:
diff --git a/tests/test_ipynb.py b/tests/test_ipynb.py
--- 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: