Keep IPython magics in annotated assignments in notebooks (#5628)

IPython transforms an assignment magic with a type annotation the same
way as a plain one, so

    files: list[str] = !ls *.csv

becomes `files: list[str] = get_ipython().getoutput('ls *.csv')`. The
MagicFinder only looked at `ast.Assign` nodes, so the magic in an
annotated assignment was never masked and Black wrote the transformed
`get_ipython()` call back into the cell. Look for magics in
`ast.AnnAssign` nodes too.

Author: Anandb71

CHANGES.md

@@ -61,6 +61,9 @@
   line, where it no longer applies (#5618)
 - Fix crash when a comment and a `# type: ignore` from the next line are joined onto one
   line and ensure the `# type: ignore` comes first (#5625)
+- Fix a Jupyter notebook magic in an annotated assignment (for example
+  `files: list[str] = !ls`) being replaced with the `get_ipython()` call that IPython
+  runs it as (#5628)
 
 ### Preview style
 

src/black/handle_ipynb_magics.py

@@ -451,18 +451,20 @@ class MagicFinder(ast.NodeVisitor):
     def __init__(self) -> None:
         self.magics: dict[int, list[OffsetAndMagic]] = collections.defaultdict(list)
 
-    def visit_Assign(self, node: ast.Assign) -> None:
+    def visit_Assign(self, node: ast.Assign | ast.AnnAssign) -> None:
         """Look for system assign magics.
 
         For example,
 
             black_version = !black --version
             env = %env var
+            files: list[str] = !ls
 
         would have been (respectively) transformed to
 
             black_version = get_ipython().getoutput('black --version')
             env = get_ipython().run_line_magic('env', 'var')
+            files: list[str] = get_ipython().getoutput('ls')
 
         and we look for instances of any of the latter.
         """
@@ -484,6 +486,10 @@ def visit_Assign(self, node: ast.Assign) -> None:
             )
         self.generic_visit(node)
 
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        """Look for system assign magics with a type annotation."""
+        self.visit_Assign(node)
+
     def visit_Expr(self, node: ast.Expr) -> None:
         """Look for magics in body of cell.
 

tests/test_ipynb.py

@@ -144,6 +144,16 @@ def test_cell_magic_noop() -> None:
             id="Assignment after non-ASCII characters",
         ),
         pytest.param("env =  %env", "env = %env", id="Assignment to magic"),
+        pytest.param(
+            "files: list[str] =!ls",
+            "files: list[str] = !ls",
+            id="Annotated system assignment",
+        ),
+        pytest.param(
+            "home: str =  %env HOME",
+            "home: str = %env HOME",
+            id="Annotated assignment to environment variable",
+        ),
     ),
 )
 def test_magic(src: str, expected: str, mode: Mode) -> None:
@@ -206,6 +216,11 @@ def test_magic_noop() -> None:
             " --format=freeze --disable-pip-version-check",
             id="Indented system assignment",
         ),
+        pytest.param(
+            "installed_packages_in_this_environment: list[str] = !pip list"
+            " --format=freeze --disable-pip-version-check",
+            id="Annotated system assignment",
+        ),
     ),
 )
 def test_long_magic_assignment_is_not_wrapped_in_parentheses(