PY-20058 Suppress E402 errors from PEP 8 inspection in IPython notebook cells

The error is suppressed only when magic commands cause it, not any
other unexpected statement preceding an import. It doesn't work for
docstrings preceded by magic commands since our lexer cannot recognize
such dangling string literals as docstrings in the first place, but it
seems out of the scope for this issue.

diff --git a/python/ipnb/src/org/jetbrains/plugins/ipnb/IpnbPep8ProblemSuppressor.java b/python/ipnb/src/org/jetbrains/plugins/ipnb/IpnbPep8ProblemSuppressor.java
index 4b68d727b76f41b1bfaaa9771c6e17015a80be54..5e152740ae7523ff86f08b269e27a55a3d16eda5 100644 (file)
--- a/python/ipnb/src/org/jetbrains/plugins/ipnb/IpnbPep8ProblemSuppressor.java
+++ b/python/ipnb/src/org/jetbrains/plugins/ipnb/IpnbPep8ProblemSuppressor.java
@@ -2,12 +2,20 @@ package org.jetbrains.plugins.ipnb;
 
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.PyEmptyExpression;
+import com.jetbrains.python.psi.PyExpressionStatement;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.PyImportStatement;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.validation.Pep8ExternalAnnotator;
 import com.jetbrains.python.validation.Pep8ProblemSuppressor;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.ipnb.psi.IpnbPyFragment;
 
+import static com.jetbrains.python.psi.PyUtil.as;
+
 /**
  * @author Mikhail Golubev
  */
@@ -16,7 +24,36 @@ public class IpnbPep8ProblemSuppressor implements Pep8ProblemSuppressor {
   public boolean isProblemSuppressed(@NotNull Pep8ExternalAnnotator.Problem problem,
                                      @NotNull PsiFile file,
                                      @Nullable PsiElement targetElement) {
-    // Ignore warnings about missing new line at the end of file inside IPython notebook cells
-    return file instanceof IpnbPyFragment && problem.getCode().equals("W292");
+    if (file instanceof IpnbPyFragment) {
+      // Ignore warnings about missing new line at the end of file
+      if (problem.getCode().equals("W292")) {
+        return true;
+      }
+
+      // Ignore warnings about imports not at the top of file if there are magic commands before
+      if (problem.getCode().equals("E402") && targetElement != null) {
+        final PyImportStatement importStatement = PsiTreeUtil.getParentOfType(targetElement, PyImportStatement.class);
+        if (importStatement != null && importStatement.getParent() == file) {
+          boolean containsMagicCommandsBefore = false;
+          PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(importStatement);
+          while (prev != null) {
+            if (prev instanceof PyEmptyExpression && prev.getText().startsWith("%")) {
+              containsMagicCommandsBefore = true;
+            }
+            else if (!isModuleLevelDocstring(prev, (PyFile)file)) {
+              return false;
+            }
+            prev = PyPsiUtils.getPrevNonWhitespaceSibling(prev);
+          }
+          return containsMagicCommandsBefore;
+        }
+      }
+    }
+    return false;
+  }
+
+  private static boolean isModuleLevelDocstring(@NotNull PsiElement element, @NotNull PyFile file) {
+    final PyExpressionStatement expressionStatement = as(element, PyExpressionStatement.class);
+    return expressionStatement != null && expressionStatement.getExpression() == file.getDocStringExpression();
   }
 }