diff --git a/python/src/META-INF/python-plugin-common.xml b/python/src/META-INF/python-plugin-common.xml
index 8cf8e8428738..ca04dcd5d281 100644
--- a/python/src/META-INF/python-plugin-common.xml
+++ b/python/src/META-INF/python-plugin-common.xml
@@ -127,6 +127,16 @@
Python
+
+ com.jetbrains.python.codeInsight.intentions.SpecifyTypeInDocstringIntention
+ Python
+
+
+
+ com.jetbrains.python.codeInsight.intentions.TypeAssertionIntention
+ Python
+
+
com.jetbrains.python.codeInsight.intentions.ImportFromToImportIntention
Python
diff --git a/python/src/com/jetbrains/python/PyBundle.properties b/python/src/com/jetbrains/python/PyBundle.properties
index fd8828138357..9f3bc849bdd6 100644
--- a/python/src/com/jetbrains/python/PyBundle.properties
+++ b/python/src/com/jetbrains/python/PyBundle.properties
@@ -182,6 +182,12 @@ INTN.transform.into.if.else.statement=Transform conditional expression into if/e
# PyConvertTripleQuotedStringIntention
INTN.doc.string.stub=Insert documentation string stub
+#SpecifyTypeInDocstringIntention
+INTN.specify.type=Specify type for reference in docstring
+
+#TypeAssertionIntention
+INTN.insert.assertion=Insert type assertion
+
# Conflict checker
CONFLICT.name.$0.obscured=Name ''{0}'' obscured by local definitions
CONFLICT.name.$0.obscured.cannot.convert=Name ''{0}'' obscured. Cannot convert.
diff --git a/python/src/com/jetbrains/python/actions/DocstringQuickFix.java b/python/src/com/jetbrains/python/actions/DocstringQuickFix.java
index 6dff34a4442f..3054dccb9981 100644
--- a/python/src/com/jetbrains/python/actions/DocstringQuickFix.java
+++ b/python/src/com/jetbrains/python/actions/DocstringQuickFix.java
@@ -8,9 +8,7 @@ import com.intellij.openapi.editor.EditorFactory;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.text.LineTokenizer;
import com.intellij.psi.PsiDocumentManager;
-import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.documentation.PyDocumentationSettings;
@@ -64,7 +62,9 @@ public class DocstringQuickFix implements LocalQuickFix {
private static Editor getEditor(Project project, PsiFile file) {
Document document = PsiDocumentManager.getInstance(project).getDocument(file);
if (document != null) {
- Editor[] editors = EditorFactory.getInstance().getEditors(document);
+ final EditorFactory instance = EditorFactory.getInstance();
+ if (instance == null) return null;
+ Editor[] editors = instance.getEditors(document);
if (editors.length > 0)
return editors[0];
}
@@ -74,14 +74,12 @@ public class DocstringQuickFix implements LocalQuickFix {
public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PyDocStringOwner.class);
if (docStringOwner == null) return;
- PyStringLiteralExpression element = docStringOwner.getDocStringExpression();
- if (element == null && myMissing == null && myUnexpected == null) {
+ PyStringLiteralExpression docStringExpression = docStringOwner.getDocStringExpression();
+ if (docStringExpression == null && myMissing == null && myUnexpected == null) {
if (docStringOwner instanceof PyFunction) {
- PsiDocumentManager.getInstance(project).getDocument(docStringOwner.getContainingFile());
PythonDocumentationProvider.inserDocStub((PyFunction)docStringOwner, project, getEditor(project, docStringOwner.getContainingFile()));
}
if (docStringOwner instanceof PyClass) {
- PsiDocumentManager.getInstance(project).getDocument(docStringOwner.getContainingFile());
PyFunction init = ((PyClass)docStringOwner).findInitOrNew(false);
if (init == null) return;
PythonDocumentationProvider.inserDocStub(init, ((PyClass)docStringOwner).getStatementList(),
@@ -89,25 +87,26 @@ public class DocstringQuickFix implements LocalQuickFix {
}
return;
}
+ if (docStringExpression == null) return;
PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
- PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(element.getProject());
- if (documentationSettings.isEpydocFormat(element.getContainingFile())) {
+ PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project);
+ if (documentationSettings.isEpydocFormat(docStringExpression.getContainingFile())) {
myPrefix = "@";
}
else {
myPrefix = ":";
}
- String replacement = element.getText();
+ String replacement = docStringExpression.getText();
if (myMissing != null) {
- replacement = createMissingReplacement(element);
+ replacement = createMissingReplacement(docStringOwner);
}
if (myUnexpected != null) {
replacement = createUnexpectedReplacement(replacement);
}
- if (!replacement.equals(element.getText())) {
+ if (!replacement.equals(docStringExpression.getText())) {
PyExpression str = elementGenerator.createDocstring(replacement).getExpression();
- element.replace(str);
+ docStringExpression.replace(str);
}
}
@@ -145,66 +144,7 @@ public class DocstringQuickFix implements LocalQuickFix {
return newText.toString();
}
- private String createMissingReplacement(PsiElement element) {
- String text = element.getText();
- String[] lines = LineTokenizer.tokenize(text, true);
- StringBuilder replacementText = new StringBuilder();
- int ind = lines.length - 1;
- if (lines.length == 1) {
- return createSingleLineReplacement(element);
- }
- for (int i = 0; i != lines.length - 1; ++i) {
- String line = lines[i];
- if (line.contains(myPrefix)) {
- ind = i;
- break;
- }
- replacementText.append(line);
- }
- addParam(replacementText, element, false);
- for (int i = ind; i != lines.length; ++i) {
- String line = lines[i];
- replacementText.append(line);
- }
- return replacementText.toString();
- }
-
- private void addParam(StringBuilder replacementText, PsiElement element, boolean addWS) {
- PyFunction fun = PsiTreeUtil.getParentOfType(element, PyFunction.class);
- PsiWhiteSpace whitespace = PsiTreeUtil.getPrevSiblingOfType(fun.getStatementList(), PsiWhiteSpace.class);
- String ws = "\n";
- if (whitespace != null) {
- String[] spaces = whitespace.getText().split("\n");
- if (spaces.length > 1) {
- ws = ws + whitespace.getText().split("\n")[1];
- }
- }
- if (replacementText.length() > 0)
- replacementText.deleteCharAt(replacementText.length() - 1);
- replacementText.append(ws);
-
- String paramText = myMissingText;
- replacementText.append(myPrefix).append("param ").append(paramText).append(": ");
- if (addWS)
- replacementText.append(ws);
- else
- replacementText.append("\n");
- }
-
- private String createSingleLineReplacement(PsiElement element) {
- String text = element.getText();
- StringBuilder replacementText = new StringBuilder();
- String closingQuotes = "";
- if (text.endsWith("'''") || text.endsWith("\"\"\"")) {
- replacementText.append(text.substring(0, text.length() - 2));
- closingQuotes = text.substring(text.length() - 3);
- }
- else {
- replacementText.append(text.substring(0, text.length()));
- closingQuotes = text.substring(text.length() - 1);
- }
- addParam(replacementText, element, true);
- replacementText.append(closingQuotes);
- return replacementText.toString();
+ private String createMissingReplacement(PyDocStringOwner docstring) {
+ return PythonDocumentationProvider.addParamToDocstring(docstring, "param", myMissingText, myPrefix).getFirst();
}
}
diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java
index ddea272dc7d7..9f35b489520e 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java
@@ -11,6 +11,7 @@ import com.jetbrains.python.PyBundle;
import com.jetbrains.python.documentation.PythonDocumentationProvider;
import com.jetbrains.python.psi.PyDocStringOwner;
import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyStatementList;
import org.jetbrains.annotations.NotNull;
/**
@@ -33,13 +34,16 @@ public class PyDocStubIntention extends BaseIntentionAction {
public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
PyFunction function = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyFunction.class);
- if (function != null) {
+ PyStatementList list = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyStatementList.class,
+ false, PyFunction.class);
+ if (function != null && list == null) {
final PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()),
PyDocStringOwner.class);
if (docStringOwner != null) {
if (docStringOwner.getDocStringExpression() != null) return false;
}
- if (function.getStatementList() != null && function.getStatementList().getStatements().length != 0)
+ final PyStatementList statementList = function.getStatementList();
+ if (statementList != null && statementList.getStatements().length != 0)
return true;
}
return false;
diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java
new file mode 100644
index 000000000000..b01b2ad6c93c
--- /dev/null
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java
@@ -0,0 +1,145 @@
+package com.jetbrains.python.codeInsight.intentions;
+
+import com.intellij.codeInsight.CodeInsightUtilBase;
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
+import com.jetbrains.python.documentation.PyDocumentationSettings;
+import com.jetbrains.python.documentation.PythonDocumentationProvider;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyReturnTypeReference;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User: ktisha
+ *
+ * Helps to specify type
+ */
+public class SpecifyTypeInDocstringIntention implements IntentionAction {
+
+ public SpecifyTypeInDocstringIntention() {
+ }
+
+ @NotNull
+ public String getText() {
+ return PyBundle.message("INTN.specify.type");
+ }
+
+ @NotNull
+ public String getFamilyName() {
+ return PyBundle.message("INTN.specify.type");
+ }
+
+ public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+ PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1),
+ PyNamedParameter.class);
+ if (problemElement == null)
+ problemElement = PsiTreeUtil.getTopmostParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1),
+ PyQualifiedExpression.class);
+ if (problemElement == null) return false;
+ if (problemElement instanceof PyQualifiedExpression) {
+ final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
+ if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) {
+ problemElement = qualifier;
+ }
+ }
+ final PyType type = problemElement.getType(TypeEvalContext.fast());
+ if (type == null || type instanceof PyReturnTypeReference) {
+ PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class);
+ PsiReference reference = problemElement.getReference();
+ if (problemElement instanceof PyQualifiedExpression) {
+ final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
+ if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) reference = qualifier.getReference();
+ }
+ if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter))
+ return true;
+ }
+ return false;
+ }
+
+ public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+ PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1), PyNamedParameter.class);
+ if (problemElement == null)
+ problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), PyExpression.class);
+ if (problemElement != null) {
+ String name = problemElement.getName();
+ PsiReference reference = problemElement.getReference();
+ if (problemElement instanceof PyQualifiedExpression) {
+ final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
+ if (qualifier != null) {
+ reference = qualifier.getReference();
+ name = qualifier.getText();
+ }
+ }
+ PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project);
+ String prefix = ":";
+ if (documentationSettings.isEpydocFormat(file)) {
+ prefix = "@";
+ }
+ PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+
+ PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class);
+
+ if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) {
+ PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression();
+ int startOffset;
+ int endOffset;
+ final Pair replacementToOffset =
+ PythonDocumentationProvider.addParamToDocstring(pyFunction, "type", name, prefix);
+
+ if (docStringExpression != null) {
+ final String typePattern = "type " + name + ":";
+ final int index = docStringExpression.getText().indexOf(typePattern);
+ if (index == -1) {
+ PyExpression str = elementGenerator.createDocstring(replacementToOffset.getFirst()).getExpression();
+ docStringExpression.replace(str);
+ startOffset = replacementToOffset.getSecond();
+ endOffset = startOffset;
+ }
+ else {
+ startOffset = index + typePattern.length() + 1;
+ endOffset = docStringExpression.getText().indexOf("\n", startOffset);
+ if (endOffset == -1) endOffset = startOffset;
+ }
+ pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
+ docStringExpression = pyFunction.getDocStringExpression();
+ }
+ else {
+ PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst());
+ final PyStatementList list = pyFunction.getStatementList();
+ if (list != null && list.getStatements().length != 0) {
+ list.addBefore(str, list.getStatements()[0]);
+ }
+
+ pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction);
+ docStringExpression = pyFunction.getDocStringExpression();
+ startOffset = replacementToOffset.getSecond();
+ endOffset = startOffset;
+ }
+ assert docStringExpression != null;
+ int textOffSet = docStringExpression.getTextOffset();
+ editor.getCaretModel().moveToOffset(textOffSet);
+
+ final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression);
+ builder.replaceRange(TextRange.create(startOffset, endOffset), PyNames.OBJECT);
+ Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+ TemplateManager.getInstance(project).startTemplate(editor, template);
+ }
+ }
+ }
+
+ public boolean startInWriteAction() {
+ return true;
+ }
+}
\ No newline at end of file
diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java
new file mode 100644
index 000000000000..8806f9a3dbbf
--- /dev/null
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java
@@ -0,0 +1,101 @@
+package com.jetbrains.python.codeInsight.intentions;
+
+import com.intellij.codeInsight.CodeInsightUtilBase;
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.codeInsight.template.*;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyNames;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.types.PyReturnTypeReference;
+import com.jetbrains.python.psi.types.PyType;
+import com.jetbrains.python.psi.types.TypeEvalContext;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * User: ktisha
+ *
+ * Helps to specify type by assertion
+ */
+public class TypeAssertionIntention implements IntentionAction {
+
+ public TypeAssertionIntention() {
+ }
+
+ @NotNull
+ public String getText() {
+ return PyBundle.message("INTN.insert.assertion");
+ }
+
+ @NotNull
+ public String getFamilyName() {
+ return PyBundle.message("INTN.insert.assertion");
+ }
+
+ public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+ PyExpression problemElement =
+ PsiTreeUtil.getTopmostParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), PyQualifiedExpression.class);
+ if (problemElement == null) return false;
+ if (problemElement instanceof PyQualifiedExpression && ((PyQualifiedExpression)problemElement).getReferencedName() == null) {
+ final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
+ if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) {
+ problemElement = qualifier;
+ }
+ }
+ final PyType type = problemElement.getType(TypeEvalContext.fast());
+ return (type == null || type instanceof PyReturnTypeReference);
+ }
+
+ public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+ PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1),
+ PyExpression.class);
+ if (problemElement != null) {
+ PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
+
+ String name = problemElement.getText();
+ if (problemElement instanceof PyQualifiedExpression) {
+ final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier();
+ if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) {
+ final String referencedName = ((PyQualifiedExpression)problemElement).getReferencedName();
+ if (referencedName == null || PyNames.GETITEM.equals(referencedName))
+ name = qualifier.getText();
+ }
+ }
+
+ final String text = "assert isinstance(" + name + ", )";
+ PyAssertStatement assertStatement = elementGenerator.createFromText(LanguageLevel.forElement(problemElement),
+ PyAssertStatement.class, text);
+
+ final PsiElement parentStatement = PsiTreeUtil.getParentOfType(problemElement, PyStatement.class);
+ if (parentStatement == null) return;
+ final PsiElement parent = parentStatement.getParent();
+ PsiElement element;
+ if (parentStatement instanceof PyAssignmentStatement &&
+ ((PyAssignmentStatement)parentStatement).getTargets()[0] == problemElement) {
+ element = parent.addAfter(assertStatement, parentStatement);
+ }
+ else {
+ element = parent.addBefore(assertStatement, parentStatement);
+ }
+
+ int textOffSet = element.getTextOffset();
+ editor.getCaretModel().moveToOffset(textOffSet);
+
+ element = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(element);
+ final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(element);
+ builder.replaceRange(TextRange.create(text.length()-1, text.length()-1), PyNames.OBJECT);
+ Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
+ TemplateManager.getInstance(project).startTemplate(editor, template);
+ }
+ }
+
+ public boolean startInWriteAction() {
+ return true;
+ }
+}
\ No newline at end of file
diff --git a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java
index 4785b7fd8d8a..8ed5fc982957 100644
--- a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java
+++ b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java
@@ -11,6 +11,8 @@ import com.intellij.openapi.project.Project;
import com.intellij.openapi.projectRoots.Sdk;
import com.intellij.openapi.roots.ProjectRootManager;
import com.intellij.openapi.ui.Messages;
+import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.text.LineTokenizer;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.*;
@@ -34,7 +36,10 @@ import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.io.IOException;
-import java.util.*;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import static com.jetbrains.python.documentation.DocumentationBuilderKit.*;
@@ -473,7 +478,7 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
if (whitespace != null) {
String[] spaces = whitespace.getText().split("\n");
if (spaces.length > 1)
- ws = ws + spaces[spaces.length-1];
+ ws += spaces[spaces.length - 1];
}
String docContent = ws + generateDocumentationContentStub(function, ws, true);
PyExpressionStatement string = elementGenerator.createDocstring("\"\"\"" + docContent + "\"\"\"");
@@ -497,7 +502,7 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
if (whitespace != null) {
String[] spaces = whitespace.getText().split("\n");
if (spaces.length > 1)
- ws = ws + whitespace.getText().split("\n")[1];
+ ws += whitespace.getText().split("\n")[1];
}
return generateDocumentationContentStub(element, ws, checkReturn);
}
@@ -545,6 +550,86 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i
return builder.toString();
}
+ public static Pair addParamToDocstring(PyDocStringOwner function, String keyword, String paramName, String prefix) {
+ final PyStringLiteralExpression docstring = function.getDocStringExpression();
+ String text;
+ if (docstring != null)
+ text = docstring.getText();
+ else
+ text = "\"\"\"\"\"\"";
+ String[] lines = LineTokenizer.tokenize(text, true);
+ StringBuilder replacementText = new StringBuilder();
+ int ind = lines.length - 1;
+ if (lines.length == 1) {
+ return createSingleLineReplacement(function, keyword, paramName, prefix);
+ }
+ for (int i = 0; i != lines.length - 1; ++i) {
+ String line = lines[i];
+ if (line.contains(prefix)) {
+ ind = i;
+ break;
+ }
+ replacementText.append(line);
+ }
+ int offset = addParamOrType(replacementText, function, false, keyword, paramName, prefix);
+ for (int i = ind; i != lines.length; ++i) {
+ String line = lines[i];
+ replacementText.append(line);
+ }
+ return new Pair(replacementText.toString(), offset);
+ }
+
+ private static int addParamOrType(StringBuilder replacementText, PyDocStringOwner function, boolean addWS, String keyword,
+ String paramName, String prefix) {
+ PsiWhiteSpace whitespace = null;
+ if (function instanceof PyFunction)
+ whitespace = PsiTreeUtil.getPrevSiblingOfType(((PyFunction)function).getStatementList(), PsiWhiteSpace.class);
+ String ws = "\n";
+ if (whitespace != null) {
+ String[] spaces = whitespace.getText().split("\n");
+ if (spaces.length > 1) {
+ ws += whitespace.getText().split("\n")[1];
+ }
+ }
+ if (replacementText.length() > 0)
+ replacementText.deleteCharAt(replacementText.length() - 1);
+ replacementText.append(ws);
+
+ replacementText.append(prefix);
+ replacementText.append(keyword);
+ replacementText.append(" ");
+ replacementText.append(paramName);
+ replacementText.append(": ");
+ int offset = replacementText.length();
+ if (addWS)
+ replacementText.append(ws);
+ else
+ replacementText.append("\n");
+ return offset;
+ }
+
+ private static Pair createSingleLineReplacement(PyDocStringOwner function, String keyword, String paramName, String prefix) {
+ String text;
+ final PyStringLiteralExpression docstring = function.getDocStringExpression();
+ if (docstring != null)
+ text = docstring.getText();
+ else
+ text = "\"\"\"\"\"\"";
+ StringBuilder replacementText = new StringBuilder();
+ String closingQuotes = "";
+ if (text.endsWith("'''") || text.endsWith("\"\"\"")) {
+ replacementText.append(text.substring(0, text.length() - 2));
+ closingQuotes = text.substring(text.length() - 3);
+ }
+ else {
+ replacementText.append(text.substring(0, text.length()));
+ closingQuotes = text.substring(text.length() - 1);
+ }
+ final int offset = addParamOrType(replacementText, function, true, keyword, paramName, prefix);
+ replacementText.append(closingQuotes);
+ return new Pair(replacementText.toString(), offset);
+ }
+
private static class RaiseVisitor extends PyRecursiveElementVisitor {
private boolean myHasRaise = false;
private boolean myHasReturn = false;
diff --git a/python/testData/intentions/afterTypeAssertion.py b/python/testData/intentions/afterTypeAssertion.py
new file mode 100644
index 000000000000..25419dd10530
--- /dev/null
+++ b/python/testData/intentions/afterTypeAssertion.py
@@ -0,0 +1,4 @@
+def foo(a, b):
+ assert isinstance(a, object)
+ a.
+ b = 1
\ No newline at end of file
diff --git a/python/testData/intentions/afterTypeInDocstring.py b/python/testData/intentions/afterTypeInDocstring.py
new file mode 100644
index 000000000000..910a0afdb11f
--- /dev/null
+++ b/python/testData/intentions/afterTypeInDocstring.py
@@ -0,0 +1,6 @@
+def foo(a, b):
+ """
+ :type a: object
+ """
+ a.
+ b = 1
\ No newline at end of file
diff --git a/python/testData/intentions/beforeTypeAssertion.py b/python/testData/intentions/beforeTypeAssertion.py
new file mode 100644
index 000000000000..d34e3f2d3c92
--- /dev/null
+++ b/python/testData/intentions/beforeTypeAssertion.py
@@ -0,0 +1,3 @@
+def foo(a, b):
+ a.
+ b = 1
\ No newline at end of file
diff --git a/python/testData/intentions/beforeTypeInDocstring.py b/python/testData/intentions/beforeTypeInDocstring.py
new file mode 100644
index 000000000000..d34e3f2d3c92
--- /dev/null
+++ b/python/testData/intentions/beforeTypeInDocstring.py
@@ -0,0 +1,3 @@
+def foo(a, b):
+ a.
+ b = 1
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/PyIntentionTest.java
index 7389a52b9c61..33a0ecb0f03c 100644
--- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java
@@ -226,6 +226,14 @@ public class PyIntentionTest extends PyTestCase {
doTest("Convert to 'import sys'");
}
+ public void testTypeInDocstring() {
+ doTest(PyBundle.message("INTN.specify.type"));
+ }
+
+ public void testTypeAssertion() {
+ doTest(PyBundle.message("INTN.insert.assertion"));
+ }
+
public void testDocStub() {
CodeInsightSettings codeInsightSettings = CodeInsightSettings.getInstance();
codeInsightSettings.JAVADOC_STUB_ON_ENTER = true;