From 358edc489f79801800a0742bfc5785e5de02a550 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 3 Jul 2018 15:38:28 +0300
Subject: [PATCH] While inferring parameter type from usages, check if it is
 compared with `None` (PY-27231)

---
 .../python/psi/impl/PyNamedParameterImpl.java | 38 ++++++++++---
 .../PyTypeCheckerInspectionTest.java          | 55 +++++++++++++++++++
 2 files changed, 86 insertions(+), 7 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
index 5c38699137c7..32e52a3eb175 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java
@@ -304,9 +304,9 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
           }
         }
         if (context.maySwitchToAST(this)) {
-          final Set<String> attributes = collectUsedAttributes(context);
-          if (!attributes.isEmpty()) {
-            return new PyStructuralType(attributes, true);
+          final PyType typeFromUsages = getTypeFromUsages(context);
+          if (typeFromUsages != null) {
+            return typeFromUsages;
           }
         }
       }
@@ -319,12 +319,15 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
     return new PyElementPresentation(this);
   }
 
-  @NotNull
-  private Set<String> collectUsedAttributes(@NotNull final TypeEvalContext context) {
-    final Set<String> result = new LinkedHashSet<>();
+  @Nullable
+  private PyType getTypeFromUsages(@NotNull TypeEvalContext context) {
+    final Set<String> usedAttributes = new LinkedHashSet<>();
+
     final ScopeOwner owner = ScopeUtil.getScopeOwner(this);
     final String name = getName();
+
     final Ref<Boolean> parameterWasReassigned = Ref.create(false);
+    final Ref<Boolean> noneComparison = Ref.create(false);
 
     if (owner != null && name != null) {
       owner.accept(new PyRecursiveElementVisitor() {
@@ -411,6 +414,21 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
           }
         }
 
+        @Override
+        public void visitPyBinaryExpression(PyBinaryExpression node) {
+          super.visitPyBinaryExpression(node);
+
+          if (noneComparison.get() || !node.isOperator(PyNames.IS) && !node.isOperator("isnot")) return;
+
+          final PyExpression lhs = node.getLeftExpression();
+          final PyExpression rhs = node.getRightExpression();
+
+          if (isReferenceToParameter(lhs) ^ isReferenceToParameter(rhs) &&
+              (lhs != null && context.getType(lhs) instanceof PyNoneType) ^ (rhs != null && context.getType(rhs) instanceof PyNoneType)) {
+            noneComparison.set(true);
+          }
+        }
+
         @Contract("null -> false")
         private boolean isReferenceToParameter(@Nullable PsiElement element) {
           if (element == null) return false;
@@ -419,7 +437,13 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
         }
       });
     }
-    return result;
+
+    if (!usedAttributes.isEmpty()) {
+      final PyStructuralType structuralType = new PyStructuralType(usedAttributes, true);
+      return noneComparison.get() ? PyUnionType.union(structuralType, PyNoneType.INSTANCE) : structuralType;
+    }
+
+    return null;
   }
 
   @NotNull
diff --git a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
index 9a6ca49b871b..5b570d152f0c 100644
--- a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java
@@ -579,6 +579,61 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-27231
+  public void testStructuralAndNone() {
+    doTestByText("def func11(value):\n" +
+                 "    if value is not None and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func12(value):\n" +
+                 "    if None is not value and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func21(value):\n" +
+                 "    if value is None and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func22(value):\n" +
+                 "    if None is value and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "func11(None)\n" +
+                 "func12(None)\n" +
+                 "func21(None)\n" +
+                 "func22(None)\n" +
+                 "\n" +
+                 "\n" +
+                 "def func31(value):\n" +
+                 "    if value and None and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func32(value):\n" +
+                 "    if value is value and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func33(value):\n" +
+                 "    if None is None and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "def func34(value):\n" +
+                 "    a = 2\n" +
+                 "    if a is a and value != 1:\n" +
+                 "        pass\n" +
+                 "\n" +
+                 "\n" +
+                 "func31(<warning descr=\"Expected type '{__ne__}', got 'None' instead\">None</warning>)\n" +
+                 "func32(<warning descr=\"Expected type '{__ne__}', got 'None' instead\">None</warning>)\n" +
+                 "func33(<warning descr=\"Expected type '{__ne__}', got 'None' instead\">None</warning>)\n" +
+                 "func34(<warning descr=\"Expected type '{__ne__}', got 'None' instead\">None</warning>)");
+  }
+
   // PY-29704
   public void testPassingAbstractMethodResult() {
     doTestByText("import abc\n" +