From 149eefa75e4c97ad64c4889423aba8a1776e2d26 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Mon, 29 Jan 2018 18:45:53 +0300
Subject: [PATCH] PY-16412 Proper use of a smart pointer in case of malformed
 annotations

Namely, it was mistakenly dereferenced before the document was
committed after modifications.

Callable types are still not rendered properly in PEP 484 compatible
format, but at least they don't lead to PIEAE anymore.
---
 .../codeInsight/intentions/PyTypeHintGenerationUtil.java  | 4 ++--
 .../annotationCallableType.py                             | 5 +++++
 .../annotationCallableType_after.py                       | 8 ++++++++
 .../intentions/PyAnnotateVariableTypeIntentionTest.java   | 4 ++++
 4 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
index ad8871681322..a2cf4877c20a 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyTypeHintGenerationUtil.java
@@ -128,10 +128,10 @@ public class PyTypeHintGenerationUtil {
       final PyAssignmentStatement assignment = (PyAssignmentStatement)target.getParent();
       final SmartPsiElementPointer<PyAssignmentStatement> pointer = manager.createSmartPsiElementPointer(assignment);
       addOrUpdateAnnotatedStatement = () -> {
-        return PyUtil.updateDocumentUnblockedAndCommitted(target, document -> {
+        PyUtil.updateDocumentUnblockedAndCommitted(target, document -> {
           document.insertString(target.getTextRange().getEndOffset(), ": " + info.getAnnotationText());
-          return pointer.getElement();
         });
+        return pointer.getElement();
       };
     }
     else {
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType.py
new file mode 100644
index 000000000000..aa6520a5a912
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType.py
@@ -0,0 +1,5 @@
+def func(x):
+    pass
+
+
+va<caret>r = func
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py
new file mode 100644
index 000000000000..74fb17c6a316
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationCallableType_after.py
@@ -0,0 +1,8 @@
+from typing import Any, Callable
+
+
+def func(x):
+    pass
+
+
+var: Callable[[Any], None] = func
diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
index dbb9fab09b28..5cc90fc19413 100644
--- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
@@ -223,6 +223,10 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
     doAnnotationTest();
   }
 
+  public void testAnnotationCallableType() {
+    doAnnotationTest();
+  }
+
   private void doAnnotationTest() {
     doTest(LanguageLevel.PYTHON36);
   }