fix assignment to env var (#680)

* fix assignment to env var

* coverage

* fix test

nbqa/handle_magics.py

@@ -56,16 +56,27 @@ def visit_Assign(self, node: ast.Assign) -> None:
         node
             Function call.
 
+        Raises
+        ------
+        AssertionError
+            If unreachable code is found.
         """
-        if (
-            isinstance(node.value, ast.Call)
-            and _is_ipython_magic(node.value.func)
-            and isinstance(node.value.func, ast.Attribute)
-            and node.value.func.attr == "getoutput"
-        ):
+        if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
+            assert isinstance(node.value.func, ast.Attribute)  # help mypy
             args = _get_node_args(node.value)
-            src = f"!{args[0]}"
-            magic_type = "line"
+            if node.value.func.attr == "getoutput":
+                src = f"!{args[0]}"
+                magic_type = "line"
+            elif node.value.func.attr == "run_line_magic":
+                src = f"%{args[0]}"
+                if args[1]:
+                    assert src is not None
+                    src += f" {args[1]}"
+                magic_type = "line"
+            else:
+                raise AssertionError(
+                    "Unreachable code: please report a bug at https://github.com/nbQA-dev/nbQA/issues"
+                )
             self.magics[node.value.lineno].append(
                 (
                     node.value.col_offset,

tests/data/env_var.ipynb

@@ -0,0 +1,40 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "skip-flake8"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "var  = %env var\n",
+    "foo = %foo"
+   ]
+  }
+ ],
+ "metadata": {
+  "anaconda-cloud": {},
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

tests/tools/test_black.py

@@ -542,3 +542,31 @@ def test_comment_after_trailing_comma(capsys: "CaptureFixture") -> None:
         "To apply these changes, remove the `--nbqa-diff` flag\n"
     )
     assert out == expected_out
+
+
+def test_assignment_to_env_var(capsys: "CaptureFixture") -> None:
+    """
+    Check that assigning to %env rountrips.
+
+    Parameters
+    ----------
+    capsys
+        Pytest fixture to capture stdout and stderr.
+    """
+    path = os.path.abspath(os.path.join("tests", "data", "env_var.ipynb"))
+
+    main(["black", path, "--nbqa-diff"])
+
+    out, _ = capsys.readouterr()
+    expected_out = (
+        "\x1b[1mCell 1\x1b[0m\n"
+        "------\n"
+        f"\x1b[1;37m--- {path}\n"
+        f"\x1b[0m\x1b[1;37m+++ {path}\n"
+        "\x1b[0m\x1b[36m@@ -1,2 +1,2 @@\n"
+        "\x1b[0m\x1b[31m-var  = %env var\n"
+        "\x1b[0m\x1b[32m+var = %env var\n"
+        "\x1b[0m\n"
+        "To apply these changes, remove the `--nbqa-diff` flag\n"
+    )
+    assert out == expected_out