perf (#619)

nbqa/handle_magics.py

@@ -2,26 +2,11 @@
 import ast
 import secrets
 from collections import defaultdict
-from typing import List, MutableMapping, Optional, Set, Tuple
+from typing import List, MutableMapping, Optional, Tuple
 
 COMMANDS_WITH_STRING_TOKEN = {"flake8"}
 
 
-class SystemAssignsFinder(ast.NodeVisitor):
-    """Find assignments of system commands."""
-
-    def __init__(self) -> None:
-        """Record where system assigns occur."""
-        self.system_assigns: Set[Tuple[int, int]] = set()
-
-    def visit_Assign(self, node: ast.Assign) -> None:  # pylint: disable=C0103
-        """Find assignments of ipython magic."""
-        if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
-            self.system_assigns.add((node.value.lineno, node.value.col_offset))
-
-        self.generic_visit(node)
-
-
 def _is_ipython_magic(node: ast.expr) -> bool:
     """Check if attribute is IPython magic."""
     return (
@@ -35,14 +20,13 @@ def _is_ipython_magic(node: ast.expr) -> bool:
 class Visitor(ast.NodeVisitor):
     """Visit cell to look for get_ipython calls."""
 
-    def __init__(self, system_assigns: Set[Tuple[int, int]]) -> None:
+    def __init__(self) -> None:
         """Magics will record where magics occur."""
         self.magics: MutableMapping[
             int, List[Tuple[int, Optional[str], Optional[str]]]
         ] = defaultdict(list)
-        self.system_assigns = system_assigns
 
-    def visit_Call(self, node: ast.Call) -> None:  # pylint: disable=C0103,R0912
+    def visit_Assign(self, node: ast.Assign) -> None:  # pylint: disable=C0103,R0912
         """
         Get source to replace ipython magic with.
 
@@ -56,10 +40,52 @@ def visit_Call(self, node: ast.Call) -> None:  # pylint: disable=C0103,R0912
         AssertionError
             Defensive check.
         """
-        if _is_ipython_magic(node.func):
-            assert isinstance(node.func, ast.Attribute)  # help mypy
+        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"
+        ):
+            args = []
+            for arg in node.value.args:
+                if isinstance(arg, ast.Str):
+                    args.append(arg.s)
+                else:
+                    raise AssertionError(
+                        "Please report a bug at https://github.com/nbQA-dev/nbQA/issues"
+                    )
+            assert (
+                args
+            ), "Please report a bug at https://github.com/nbQA-dev/nbQA/issues"
+            src = f"!{args[0]}"
+            magic_type = "line"
+            self.magics[node.value.lineno].append(
+                (
+                    node.value.col_offset,
+                    src,
+                    magic_type,
+                )
+            )
+        self.generic_visit(node)
+
+    def visit_Expr(self, node: ast.Expr) -> None:  # pylint: disable=C0103,R0912
+        """
+        Get source to replace ipython magic with.
+
+        Parameters
+        ----------
+        node
+            Function call.
+
+        Raises
+        ------
+        AssertionError
+            Defensive check.
+        """
+        if isinstance(node.value, ast.Call) and _is_ipython_magic(node.value.func):
+            assert isinstance(node.value.func, ast.Attribute)  # help mypy
             args = []
-            for arg in node.args:
+            for arg in node.value.args:
                 if isinstance(arg, ast.Str):
                     args.append(arg.s)
                 else:
@@ -70,13 +96,13 @@ def visit_Call(self, node: ast.Call) -> None:  # pylint: disable=C0103,R0912
                 args
             ), "Please report a bug at https://github.com/nbQA-dev/nbQA/issues"
             magic_type: Optional[str] = None
-            if node.func.attr == "run_cell_magic":
+            if node.value.func.attr == "run_cell_magic":
                 src: Optional[str] = f"%%{args[0]}"
                 if args[1]:
                     assert src is not None
                     src += f" {args[1]}"
                 magic_type = "cell"
-            elif node.func.attr == "run_line_magic":
+            elif node.value.func.attr == "run_line_magic":
                 if args[0] == "pinfo":
                     src = f"{args[1]}?"
                 elif args[0] == "pinfo2":
@@ -87,20 +113,17 @@ def visit_Call(self, node: ast.Call) -> None:  # pylint: disable=C0103,R0912
                         assert src is not None
                         src += f" {args[1]}"
                 magic_type = "line"
-            elif node.func.attr == "system":
+            elif node.value.func.attr == "system":
                 src = f"!{args[0]}"
                 magic_type = "line"
-            elif node.func.attr == "getoutput":
-                if (node.lineno, node.col_offset) in self.system_assigns:
-                    src = f"!{args[0]}"
-                else:
-                    src = f"!!{args[0]}"
+            elif node.value.func.attr == "getoutput":
+                src = f"!!{args[0]}"
                 magic_type = "line"
             else:
                 src = None
-            self.magics[node.lineno].append(
+            self.magics[node.value.lineno].append(
                 (
-                    node.col_offset,
+                    node.value.col_offset,
                     src,
                     magic_type,
                 )

nbqa/save_source.py

@@ -21,12 +21,7 @@
 import tokenize_rt
 from IPython.core.inputtransformer2 import TransformerManager
 
-from nbqa.handle_magics import (
-    CellMagicFinder,
-    MagicHandler,
-    SystemAssignsFinder,
-    Visitor,
-)
+from nbqa.handle_magics import CellMagicFinder, MagicHandler, Visitor
 from nbqa.notebook_info import NotebookInfo
 from nbqa.path_utils import remove_prefix
 
@@ -72,9 +67,7 @@ def _process_source(
         handler = MagicHandler(source, command, magic_type=None)
         magic_substitutions.append(handler)
         return handler.replacement
-    system_assigns_finder = SystemAssignsFinder()
-    system_assigns_finder.visit(tree)
-    visitor = Visitor(system_assigns_finder.system_assigns)
+    visitor = Visitor()
     visitor.visit(tree)
     new_src = []
     for i, line in enumerate(body.splitlines(), start=1):
@@ -267,7 +260,7 @@ def _should_ignore_code_cell(
         # Deal with this below
         pass
     else:
-        # No need to ignore
+        # Syntax is fine, no need to ignore
         return False
 
     cell_magic_finder = CellMagicFinder()