From 7d7ead47a35a06137c74e525ba0076a6670b6de2 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 13 Sep 2016 16:37:34 +0300
Subject: [PATCH] Improve "is None" and "is not None" processing in
 PyTypeAssertionEvaluator

---
 .../controlflow/PyTypeAssertionEvaluator.java | 48 +++++++++++++------
 .../com/jetbrains/python/PyTypeTest.java      | 40 ++++++++++++++++
 2 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/controlflow/PyTypeAssertionEvaluator.java b/python/src/com/jetbrains/python/codeInsight/controlflow/PyTypeAssertionEvaluator.java
index 2bce83722acf..fb49bd213239 100644
--- a/python/src/com/jetbrains/python/codeInsight/controlflow/PyTypeAssertionEvaluator.java
+++ b/python/src/com/jetbrains/python/codeInsight/controlflow/PyTypeAssertionEvaluator.java
@@ -115,21 +115,39 @@ public class PyTypeAssertionEvaluator extends PyRecursiveElementVisitor {
     final PyExpression lhs = node.getLeftExpression();
     final PyExpression rhs = node.getRightExpression();
 
-    if (node.isOperator("isnot")) {
-      if (lhs instanceof PyReferenceExpression && rhs instanceof PyReferenceExpression) {
-        final PyReferenceExpression target = (PyReferenceExpression)lhs;
-        if (PyNames.NONE.equals(rhs.getName())) {
-          final boolean positive = myPositive;
-          pushAssertion(target, new InstructionTypeCallback() {
-            @Override
-            public PyType getType(TypeEvalContext context, @Nullable PsiElement anchor) {
-              final List<PyType> types = new ArrayList<>();
-              types.add(PyNoneType.INSTANCE);
-              return createAssertionType(context.getType(target), types, !positive, context);
-            }
-          });
-          return;
-        }
+    if (lhs instanceof PyReferenceExpression && rhs instanceof PyReferenceExpression) {
+      final boolean leftIsNone = PyNames.NONE.equals(lhs.getName());
+      final boolean rightIsNone = PyNames.NONE.equals(rhs.getName());
+
+      if (leftIsNone && rightIsNone) {
+        return;
+      }
+
+      final PyReferenceExpression target = (PyReferenceExpression)(rightIsNone ? lhs : rhs);
+      final boolean positive = myPositive;
+
+      if (node.isOperator(PyNames.IS)) {
+        pushAssertion(target, new InstructionTypeCallback() {
+          @Override
+          public PyType getType(TypeEvalContext context, @Nullable PsiElement anchor) {
+            final List<PyType> types = new ArrayList<>();
+            types.add(PyNoneType.INSTANCE);
+            return createAssertionType(context.getType(target), types, positive, context);
+          }
+        });
+        return;
+      }
+
+      if (node.isOperator("isnot")) {
+        pushAssertion(target, new InstructionTypeCallback() {
+          @Override
+          public PyType getType(TypeEvalContext context, @Nullable PsiElement anchor) {
+            final List<PyType> types = new ArrayList<>();
+            types.add(PyNoneType.INSTANCE);
+            return createAssertionType(context.getType(target), types, !positive, context);
+          }
+        });
+        return;
       }
     }
 
diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java
index 8e4ccff58849..a8d51774e6cf 100644
--- a/python/testSrc/com/jetbrains/python/PyTypeTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java
@@ -1181,6 +1181,46 @@ public class PyTypeTest extends PyTestCase {
            "    print(expr)");
   }
 
+  public void testIsNotNone() {
+    doTest("int",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if x is not None:\n" +
+           "        expr = x\n");
+
+    doTest("int",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if None is not x:\n" +
+           "        expr = x\n");
+
+    doTest("int",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if not x is None:\n" +
+           "        expr = x\n");
+
+    doTest("int",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if not None is x:\n" +
+           "        expr = x\n");
+  }
+
+  public void testIsNone() {
+    doTest("None",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if x is None:\n" +
+           "        expr = x\n");
+
+    doTest("None",
+           "def test_1(self, c):\n" +
+           "    x = 1 if c else None\n" +
+           "    if None is x:\n" +
+           "        expr = x\n");
+  }
+
   private static List<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
     return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
                             TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());