From 15e6cba29a7caeb227f396b3e675496fb0235cfb Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 18 Dec 2018 21:43:52 +0300
Subject: [PATCH] Add an ability to enable mandatory encoding inspection on
 Python 3+ (PY-33087)

---
 .../inspections/PyMandatoryEncodingInspection.java  | 11 +++++++++--
 .../test.py                                         |  0
 .../inspections/AddEncodingInEnabledPy3/test.py     |  1 +
 .../com/jetbrains/python/PythonInspectionsTest.java | 13 +++++++++++--
 4 files changed, 21 insertions(+), 4 deletions(-)
 rename python/testData/inspections/{AddEncodingInPy3 => AddEncodingInDisabledPy3}/test.py (100%)
 create mode 100644 python/testData/inspections/AddEncodingInEnabledPy3/test.py

diff --git a/python/src/com/jetbrains/python/inspections/PyMandatoryEncodingInspection.java b/python/src/com/jetbrains/python/inspections/PyMandatoryEncodingInspection.java
index e1dc883fa479..4e1906a37632 100644
--- a/python/src/com/jetbrains/python/inspections/PyMandatoryEncodingInspection.java
+++ b/python/src/com/jetbrains/python/inspections/PyMandatoryEncodingInspection.java
@@ -19,6 +19,7 @@ import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.codeInspection.ProblemsHolder;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElementVisitor;
+import com.intellij.util.ui.CheckBox;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.inspections.quickfix.AddEncodingQuickFix;
@@ -29,6 +30,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import javax.swing.*;
+import java.awt.*;
 import java.awt.event.ActionEvent;
 import java.awt.event.ActionListener;
 
@@ -63,7 +65,7 @@ public class PyMandatoryEncodingInspection extends PyInspection {
 
     @Override
     public void visitPyFile(PyFile node) {
-      if (!LanguageLevel.forElement(node).isPython2()) return;
+      if (!(myAllPythons || LanguageLevel.forElement(node).isPython2())) return;
 
       final String charsetString = PythonFileType.getCharsetFromEncodingDeclaration(node);
       if (charsetString == null) {
@@ -78,6 +80,7 @@ public class PyMandatoryEncodingInspection extends PyInspection {
 
   public String myDefaultEncoding = "utf-8";
   public int myEncodingFormatIndex = 0;
+  public boolean myAllPythons = false;
 
   @Override
   public JComponent createOptionsPanel() {
@@ -103,6 +106,10 @@ public class PyMandatoryEncodingInspection extends PyInspection {
       }
     });
 
-    return PyEncodingUtil.createEncodingOptionsPanel(defaultEncoding, encodingFormat);
+    final JPanel panel = new JPanel(new BorderLayout());
+    panel.add(new CheckBox("Enable in Python 3+", this, "myAllPythons"), BorderLayout.NORTH);
+    panel.add(PyEncodingUtil.createEncodingOptionsPanel(defaultEncoding, encodingFormat), BorderLayout.CENTER);
+
+    return panel;
   }
 }
diff --git a/python/testData/inspections/AddEncodingInPy3/test.py b/python/testData/inspections/AddEncodingInDisabledPy3/test.py
similarity index 100%
rename from python/testData/inspections/AddEncodingInPy3/test.py
rename to python/testData/inspections/AddEncodingInDisabledPy3/test.py
diff --git a/python/testData/inspections/AddEncodingInEnabledPy3/test.py b/python/testData/inspections/AddEncodingInEnabledPy3/test.py
new file mode 100644
index 000000000000..76ef230348eb
--- /dev/null
+++ b/python/testData/inspections/AddEncodingInEnabledPy3/test.py
@@ -0,0 +1 @@
+<warning descr="No encoding specified for file"></warning>from typing import Any
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java b/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
index b8d273c9e98c..47ffa4857604 100644
--- a/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
@@ -271,7 +271,16 @@ public class PythonInspectionsTest extends PyTestCase {
   }
 
   // PY-32364
-  public void testAddEncodingInPy3() {
-    doHighlightingTest(PyMandatoryEncodingInspection.class, LanguageLevel.PYTHON34);
+  public void testAddEncodingInDisabledPy3() {
+    final PyMandatoryEncodingInspection inspection = new PyMandatoryEncodingInspection();
+    inspection.myAllPythons = false;
+    doHighlightingTest(inspection, LanguageLevel.PYTHON34);
+  }
+
+  // PY-32364
+  public void testAddEncodingInEnabledPy3() {
+    final PyMandatoryEncodingInspection inspection = new PyMandatoryEncodingInspection();
+    inspection.myAllPythons = true;
+    doHighlightingTest(inspection, LanguageLevel.PYTHON34);
   }
 }