From c832e56c07b83d63055f2c3243ffec2d3ec0ff81 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Wed, 1 Apr 2015 16:06:32 +0300
Subject: [PATCH] Infer 'x: Optional[T]' from 'x: T = None' according to PEP
 484 (PY-15206)

---
 .../codeInsight/PyTypingTypeProvider.java     | 22 ++++++++++++++++---
 .../com/jetbrains/python/PyTypingTest.java    |  6 +++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
index dc4f88e1ebdd..2a2f83eebefb 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyTypingTypeProvider.java
@@ -51,6 +51,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     .add("typing.Protocol")
     .build();
 
+  @Nullable
   public Ref<PyType> getParameterType(@NotNull PyNamedParameter param, @NotNull PyFunction func, @NotNull TypeEvalContext context) {
     final PyAnnotation annotation = param.getAnnotation();
     if (annotation != null) {
@@ -59,7 +60,8 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
       if (value != null) {
         final PyType type = getTypingType(value, context);
         if (type != null) {
-          return Ref.create(type);
+          final PyType optionalType = getOptionalTypeFromDefaultNone(param, type, context);
+          return Ref.create(optionalType != null ? optionalType : type);
         }
       }
     }
@@ -92,6 +94,20 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     return type instanceof PyClassType && "typing.Any".equals(((PyClassType)type).getPyClass().getQualifiedName());
   }
 
+  @Nullable
+  private static PyType getOptionalTypeFromDefaultNone(@NotNull PyNamedParameter param,
+                                                       @NotNull PyType type,
+                                                       @NotNull TypeEvalContext context) {
+    final PyExpression defaultValue = param.getDefaultValue();
+    if (defaultValue != null) {
+      final PyType defaultType = context.getType(defaultValue);
+      if (defaultType instanceof PyNoneType) {
+        return PyUnionType.union(type, defaultType);
+      }
+    }
+    return null;
+  }
+
   @Nullable
   private static PyType getGenericConstructorType(@NotNull PyFunction function, @NotNull TypeEvalContext context) {
     if (PyUtil.isInit(function)) {
@@ -152,7 +168,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
     if (unionType != null) {
       return unionType;
     }
-    final Ref<PyType> optionalType = getOptionalType(expression, context);
+    final Ref<PyType> optionalType = getOptionalTypeFromDefaultNone(expression, context);
     if (optionalType != null) {
       return optionalType.get();
     }
@@ -203,7 +219,7 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
   }
 
   @Nullable
-  private static Ref<PyType> getOptionalType(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
+  private static Ref<PyType> getOptionalTypeFromDefaultNone(@NotNull PyExpression expression, @NotNull TypeEvalContext context) {
     if (expression instanceof PySubscriptionExpression) {
       final PySubscriptionExpression subscriptionExpr = (PySubscriptionExpression)expression;
       final PyExpression operand = subscriptionExpr.getOperand();
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 959162babef6..3733dff1ea19 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -257,6 +257,12 @@ public class PyTypingTest extends PyTestCase {
            "    pass\n");
   }
 
+  public void testOptionalFromDefaultNone() {
+    doTest("Optional[int]",
+           "def foo(expr: int = None):\n" +
+           "    pass\n");
+  }
+
   private void doTest(@NotNull String expectedType, @NotNull String text) {
     myFixture.copyDirectoryToProject("typing", "");
     myFixture.configureByText(PythonFileType.INSTANCE, text);