PY-77065 Use the actual return type for "Specify return type in docstring"

Merge-request: IJ-MR-149121
Merged-by: Aleksandr Govenko <aleksandr.govenko@jetbrains.com>

GitOrigin-RevId: b71ea8e26e2253196dede08648b67f0a2ba46408

python/python-psi-impl/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java

@@ -8,6 +8,7 @@ import com.intellij.codeInsight.template.TemplateBuilderFactory;
 import com.intellij.openapi.application.WriteAction;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiFile;
@@ -17,12 +18,16 @@ import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.codeInsight.intentions.PyTypeHintGenerationUtil.AnnotationInfo;
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
 import com.jetbrains.python.debugger.PySignature;
 import com.jetbrains.python.debugger.PySignatureCacheManager;
 import com.jetbrains.python.documentation.docstrings.DocStringFormat;
 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.toolbox.Substring;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -67,21 +72,14 @@ public final class SpecifyTypeInDocstringIntention extends TypeIntention {
     if (!PyGenerateDocstringIntention.ensureNotPlainDocstringFormat(pyFunction)) return;
 
     final PyDocstringGenerator docstringGenerator = PyDocstringGenerator.forDocStringOwner(pyFunction);
-    String type = PyNames.OBJECT;
+
     if (param != null) {
-      final String paramName = StringUtil.notNullize(param.getName());
-      final PySignature signature = PySignatureCacheManager.getInstance(pyFunction.getProject()).findSignature(pyFunction);
-      if (signature != null) {
-        type = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(paramName), type);
-      }
+      String type = SpecifyTypeInPy3AnnotationsIntention.parameterType(param);
       docstringGenerator.withParamTypedByName(param, type);
     }
     else {
-      final PySignature signature = PySignatureCacheManager.getInstance(pyFunction.getProject()).findSignature(pyFunction);
-      if (signature != null) {
-        type = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), type);
-      }
-      docstringGenerator.withReturnValue(type);
+      AnnotationInfo info = SpecifyTypeInPy3AnnotationsIntention.returnType(pyFunction);
+      docstringGenerator.withReturnValue(info.getAnnotationText());
     }
 
     WriteAction.run(() -> {

python/testData/intentions/googleReturnSectionAfterKeywords_after.py

@@ -4,5 +4,5 @@ def f(**kwargs):
         foo: bar
 
     Returns:
-        object: 
+        None: 
     """

python/testData/intentions/googleReturnSectionAfterYields_after.py

@@ -4,7 +4,7 @@ def f():
         int: meaning of life, universe and everything
 
     Returns:
-        object: 
+        Generator[int, Any, None]: 
         
     Example:
         print(next(f))

python/testData/intentions/googleReturnSectionBeforeRaises_after.py

@@ -1,7 +1,7 @@
 def f():
     """
     Returns:
-        object: 
+        Any: 
 
     Raises:
         RuntimeException

python/testData/intentions/returnTypeInDocStringDifferentIndentationSize_after.py

@@ -1,7 +1,7 @@
 def func(x):
   """
 
-  :rtype: object
+  :rtype: int
   """
   print(x)
   return 1

python/testData/intentions/returnTypeInDocstring_after.py

@@ -1,7 +1,7 @@
 def f(x):
     """
 
-    :rtype: object
+    :rtype: None
     """
     pass
 

python/testData/intentions/returnTypeInEmptyGoogleDocString_after.py

@@ -1,5 +1,5 @@
 def f(x, y):
     """
     Returns:
-        object: 
+        None: 
     """

python/testData/intentions/returnTypeInEmptyNumpyDocString_after.py

@@ -2,5 +2,5 @@ def f(x, y):
     """
     Returns
     -------
-    object
+    None
     """

python/testData/intentions/returnTypeInGoogleDocStringEmptyReturnSection_after.py

@@ -3,5 +3,5 @@ def f(x, y):
     Summary.
     
     Returns:
-        object: 
+        None: 
     """

python/testData/intentions/returnTypeInNewGoogleDocString_after.py

@@ -2,6 +2,6 @@ def f(x):
     """
 
     Returns:
-        object: 
+        int: 
     """
     return 42

python/testData/intentions/returnTypeInNewNumpyDocString_after.py

@@ -3,6 +3,6 @@ def f(x):
 
     Returns
     -------
-    object
+    int
     """
     return 42

python/testData/intentions/returnTypeInNumpyDocStringEmptyReturnSection_after.py

@@ -4,5 +4,5 @@ def f(x, y):
     
     Returns
     -------
-    object
+    None
     """

python/testData/intentions/typeInDocstring1_after.py

@@ -1,6 +1,6 @@
 def foo(a, b):
   """
 
-  :rtype: object
+  :rtype: None
   """
   b = 1

python/testData/intentions/typeInDocstring2_after.py

@@ -1,7 +1,7 @@
 def func1(x):
   """
 
-  :rtype: object
+  :rtype: Any
   """
   return x
 

python/testData/intentions/typeInDocstringAtTheEndOfFunction_after.py

@@ -1,7 +1,7 @@
 def func2(x):
     """
 
-    :rtype: object
+    :rtype: Any
     """
     y = func1(x.keys())
     return y.startswith('foo')