PY-18254 PyFunction#getTypeComment returns null if comment doesn't have "type:" prefix

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
index 7f315021a6facc186f12abeda180511bc05f46d2..76c7e7ddb7ca2e26099567e24e2e0171a9e6df60 100644 (file)
--- a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
@@ -104,11 +104,10 @@ public interface PyFunction extends PsiNamedElement, StubBasedPsiElement<PyFunct
 
   /**
    * Returns a comment that immediately follows function header and precedes any child statement (including docstring).
-   * It's supposed to start with conventional "type:" prefix and contain function type in PEP 484 compatible format
+   * It must start with conventional "type:" prefix and contain function type in PEP 484 compatible format
    * (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code).
    * <p/>
-   * Use {@link #getTypeCommentAnnotation()} to make sure that the comment indeed starts with "type:" and get its content
-   * with the prefix stripped.
+   * Use {@link #getTypeCommentAnnotation()} to get its content with the prefix stripped accessing either stubs or AST.
    *
    * @see #getTypeCommentAnnotation()
    */

diff --git a/python/src/com/jetbrains/python/codeInsight/PyFunctionTypeCommentReferenceContributor.java b/python/src/com/jetbrains/python/codeInsight/PyFunctionTypeCommentReferenceContributor.java
index 06bcd88dc095383ede56ea32677ffd3bfc04abbb..63661ead56d93d1ea3398d45bc72c11e76a4bf61 100644 (file)
--- a/python/src/com/jetbrains/python/codeInsight/PyFunctionTypeCommentReferenceContributor.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyFunctionTypeCommentReferenceContributor.java
@@ -46,10 +46,7 @@ public class PyFunctionTypeCommentReferenceContributor extends PsiReferenceContr
       @Override
       public boolean accepts(@NotNull PsiComment comment, ProcessingContext context) {
         final PyFunction func = PsiTreeUtil.getParentOfType(comment, PyFunction.class);
-        if (func != null && func.getTypeComment() == comment) {
-          return PyTypingTypeProvider.getTypeCommentValue(comment.getText()) != null;
-        }
-        return false;
+        return func != null && func.getTypeComment() == comment;
       }
     });
   

diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index d027aee3cf9f83dd35079cc0f7274cd41f7ffc5b..2554a03ef4a252d7f24fcdf81e6ab32b91c887e3 100644 (file)
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -583,7 +583,10 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
   public PsiComment getTypeComment() {
     final PyStatementList statements = getStatementList();
     if (statements.getStatements().length != 0) {
-      return as(statements.getFirstChild(), PsiComment.class);
+      final PsiComment comment = as(statements.getFirstChild(), PsiComment.class);
+      if (comment != null && PyTypingTypeProvider.getTypeCommentValue(comment.getText()) != null) {
+        return comment;
+      }
     }
     return null;
   }