Generate type annotation in comment for Python 2.x

python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java

@@ -16,13 +16,18 @@
 package com.jetbrains.python.codeInsight.intentions;
 
 import com.google.common.base.Function;
+import com.google.common.collect.Lists;
 import com.intellij.codeInsight.CodeInsightUtilCore;
 import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.fileEditor.OpenFileDescriptor;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.IncorrectOperationException;
@@ -31,8 +36,9 @@ import com.jetbrains.python.documentation.doctest.PyDocstringFile;
 import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 
-import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateParameter;
-import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateReturnType;
+import java.util.List;
+
+import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.*;
 import static com.jetbrains.python.codeInsight.intentions.TypeIntention.getCallable;
 import static com.jetbrains.python.codeInsight.intentions.TypeIntention.resolvesToFunction;
 
@@ -55,8 +61,6 @@ public class PyAnnotateTypesIntention implements IntentionAction {
   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
     if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false;
 
-    if (!LanguageLevel.forElement(file).isPy3K()) return false;
-
     updateText();
 
     final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
@@ -78,7 +82,106 @@ public class PyAnnotateTypesIntention implements IntentionAction {
     final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
     final PyCallable callable = getCallable(elementAt);
 
+    if (isPy3k(file)) {
+      generatePy3kTypeAnnotations(project, editor, elementAt, callable);
+    }
+    else {
+      if (callable instanceof PyFunction) {
+        generateTypeCommentAnnotations(project, editor, elementAt, (PyFunction)callable);
+      }
+    }
+  }
+
+  private static void generateTypeCommentAnnotations(Project project, Editor editor, PsiElement at, PyFunction function) {
+
+    StringBuilder replacementTextBuilder = new StringBuilder("# type: (");
+
+    PyParameter[] params = function.getParameterList().getParameters();
+
+    List<Pair<Integer, String>> templates = Lists.newArrayList();
+
+    for (int i = 0; i < params.length; i++) {
+      String type = parameterType(params[i]);
+
+      templates.add(Pair.create(replacementTextBuilder.length(), type));
+
+      replacementTextBuilder.append(type);
+
+      if (i < params.length - 1) {
+        replacementTextBuilder.append(", ");
+      }
+    }
+
+    replacementTextBuilder.append(") -> ");
+
+    String returnType = returnType(function);
+    templates.add(Pair.create(replacementTextBuilder.length(), returnType));
+
+    replacementTextBuilder.append(returnType);
+
+    final PyStatementList statements = function.getStatementList();
+    final String indentation = PyIndentUtil.getExpectedElementIndent(statements);
+    replacementTextBuilder.insert(0, indentation);
+    replacementTextBuilder.insert(0, "\n");
 
+
+    final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
+    final Document document = manager.getDocument(function.getContainingFile());
+
+    if (document != null) {
+      final PsiElement beforeStatements = statements.getPrevSibling();
+      int offset = beforeStatements.getTextRange().getStartOffset();
+      if (":".equals(beforeStatements.getText())) {
+        offset += 1;
+      }
+      try {
+        document.insertString(offset, replacementTextBuilder.toString());
+      }
+      finally {
+        manager.commitDocument(document);
+      }
+
+
+      function = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function);
+      
+      if (function != null) {
+        final TemplateBuilder builder =
+          TemplateBuilderFactory.getInstance().createTemplateBuilder(function);
+
+        for (Pair<Integer, String> template : templates) {
+          builder.replaceRange(TextRange.from(
+            offset - function.getTextRange().getStartOffset() + replacementTextBuilder.toString().indexOf('#') + template.first,
+            template.second.length()), template.second);
+        }
+
+        startTemplate(project, function, builder);
+      }
+    }
+  }
+
+  private static void startTemplate(Project project, PyCallable callable, TemplateBuilder builder) {
+    final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+    ;
+
+    int offset = callable.getTextRange().getStartOffset();
+
+    final OpenFileDescriptor descriptor = new OpenFileDescriptor(
+      project,
+      callable.getContainingFile().getVirtualFile(),
+      offset
+    );
+    final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
+    if (targetEditor != null) {
+      targetEditor.getCaretModel().moveToOffset(offset);
+      TemplateManager.getInstance(project).startTemplate(targetEditor, template);
+    }
+  }
+
+  private static boolean isPy3k(PsiFile file) {
+    return LanguageLevel.forElement(file).isPy3K();
+  }
+
+  private static void generatePy3kTypeAnnotations(@NotNull Project project, Editor editor, PsiElement elementAt, PyCallable callable) {
     final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(callable);
 
     PyExpression returnType = annotateReturnType(project, editor.getDocument(), elementAt, false);
@@ -88,7 +191,7 @@ public class PyAnnotateTypesIntention implements IntentionAction {
     }
 
     if (callable instanceof PyFunction) {
-      PyFunction function = (PyFunction) callable;
+      PyFunction function = (PyFunction)callable;
       PyParameter[] params = function.getParameterList().getParameters();
 
       for (int i = params.length - 1; i >= 0; i--) {
@@ -112,20 +215,7 @@ public class PyAnnotateTypesIntention implements IntentionAction {
       }
     }
     if (callable != null) {
-      final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-
-      int offset = callable.getTextRange().getStartOffset();
-
-      final OpenFileDescriptor descriptor = new OpenFileDescriptor(
-        project,
-        callable.getContainingFile().getVirtualFile(),
-        offset
-      );
-      final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
-      if (targetEditor != null) {
-        targetEditor.getCaretModel().moveToOffset(offset);
-        TemplateManager.getInstance(project).startTemplate(targetEditor, template);
-      }
+      startTemplate(project, callable, builder);
     }
   }
 

python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java

@@ -91,16 +91,9 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
 
     final String defaultParamText = defaultParamValue == null ? null : defaultParamValue.getText();
 
-    String paramType = PyNames.OBJECT;
+    String paramType = parameterType(parameter);
 
-    PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class);
-    if (function != null) {
-      final PySignature signature = PySignatureCacheManager.getInstance(parameter.getProject()).findSignature(
-        function);
-      if (signature != null) {
-        paramType = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(paramName), paramType);
-      }
-    }
+   
 
     final PyNamedParameter namedParameter = elementGenerator.createParameter(paramName, defaultParamText, paramType,
                                                                              LanguageLevel.forElement(parameter));
@@ -124,18 +117,39 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
     return parameter;
   }
 
-  public static PyExpression annotateReturnType(Project project, Document document, PsiElement resolved, boolean createTemplate) {
-    PyCallable callable = getCallable(resolved);
+  static String parameterType(PyParameter parameter) {
+    String paramType = PyNames.OBJECT;
+    
+    PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class);
+    if (function != null) {
+      final PySignature signature = PySignatureCacheManager.getInstance(parameter.getProject()).findSignature(
+        function);
+      String parameterName = parameter.getName();
+      if (signature != null && parameterName != null) {
+        paramType = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(parameterName), paramType);
+      }
+    }
+    return paramType;
+  }
+
 
+  static String returnType(@NotNull PyFunction function) {
     String returnType = PyNames.OBJECT;
+    final PySignature signature = PySignatureCacheManager.getInstance(function.getProject()).findSignature(
+      function);
+    if (signature != null) {
+      returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType);
+    }
+    return returnType;
+  }
+
+  public static PyExpression annotateReturnType(Project project, Document document, PsiElement resolved, boolean createTemplate) {
+    PyCallable callable = getCallable(resolved);
 
     if (callable instanceof PyFunction) {
       PyFunction function = (PyFunction)callable;
-      final PySignature signature = PySignatureCacheManager.getInstance(project).findSignature(
-        function);
-      if (signature != null) {
-        returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType);
-      }
+
+      String returnType = returnType(function);
 
       final String annotationText = " -> " + returnType;
 
@@ -143,20 +157,22 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
       assert prevElem != null;
 
       final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
-      Document documentWithCollable = manager.getDocument(callable.getContainingFile());
-      try {
-        final TextRange range = prevElem.getTextRange();
-        manager.doPostponedOperationsAndUnblockDocument(documentWithCollable);
-        if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
-          documentWithCollable.insertString(range.getStartOffset(), annotationText);
+      Document documentWithCallable = manager.getDocument(callable.getContainingFile());
+      if (documentWithCallable != null) {
+        try {
+          final TextRange range = prevElem.getTextRange();
+          manager.doPostponedOperationsAndUnblockDocument(documentWithCallable);
+          if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
+            documentWithCallable.insertString(range.getStartOffset(), annotationText);
+          }
+          else {
+            documentWithCallable.insertString(range.getEndOffset(), annotationText + ":");
+          }
         }
-        else {
-          documentWithCollable.insertString(range.getEndOffset(), annotationText + ":");
+        finally {
+          manager.commitDocument(documentWithCallable);
         }
       }
-      finally {
-        manager.commitDocument(documentWithCollable);
-      }
 
 
       callable = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(callable);

python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment.py

+def fo<caret>o(x, y):
+    pass
\ No newline at end of file

python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment_after.py

+def foo(x, y):
+    # type: (object, object) -> object
+    pass
\ No newline at end of file

python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java

@@ -42,6 +42,11 @@ public class PyAnnotateTypesIntentionTest extends PyIntentionTestCase {
     }
   }
 
+  public void testTypeComment() {
+    doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON27);
+  }
+  
+  
   private void doTest() {
     doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON30);
   }