diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java index 3f28b98312ff..86408d90ff2b 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java +++ b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java @@ -75,6 +75,9 @@ public abstract class PyElementGenerator { @NotNull public abstract T createFromText(LanguageLevel langLevel, Class aClass, final String text, final int[] path); + public abstract PyNamedParameter createParameter(@NotNull String name, @Nullable String defaultValue, @Nullable String annotation, + @NotNull LanguageLevel level); + public abstract PyNamedParameter createParameter(@NotNull String name); public abstract PyKeywordArgument createKeywordArgument(LanguageLevel languageLevel, String keyword, String value); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index ba7ddff00c1f..82378cbd170f 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -23,6 +23,7 @@ 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; +import org.jetbrains.annotations.Nullable; /** * User: ktisha @@ -46,9 +47,11 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { if (!LanguageLevel.forElement(file).isPy3K()) return false; - PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1); + int offset = editor.getCaretModel().getOffset(); + PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) - elementAt = file.findElementAt(editor.getCaretModel().getOffset()); + elementAt = file.findElementAt(offset); PyCallExpression callExpression = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class); if (callExpression != null && callExpression.resolveCalleeFunction(PyResolveContext.defaultContext()) != null) { @@ -98,8 +101,7 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { parameter = (PyParameter)problemElement; else if (resolvedReference instanceof PyParameter) parameter = (PyParameter)resolvedReference; - if (parameter instanceof PyNamedParameter && (((PyNamedParameter)parameter).getAnnotation() != null || - parameter.getDefaultValue() != null)) return false; + if (parameter instanceof PyNamedParameter && (((PyNamedParameter)parameter).getAnnotation() != null)) return false; if (parameter != null) return true; else { @@ -146,74 +148,93 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { } } - PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + PyExpression defaultParamValue = null; + if (parameter instanceof PyNamedParameter) + defaultParamValue = parameter.getDefaultValue(); + if (parameter != null && name != null) { - final PyFunction function = - elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class, - "def foo(" + name + ": object):\n\tpass"); - final PyNamedParameter namedParameter = function.getParameterList().findParameterByName(name); - assert namedParameter != null; - parameter = (PyParameter)parameter.replace(namedParameter); - parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter); - editor.getCaretModel().moveToOffset(parameter.getTextOffset()); - - final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter); - builder.replaceRange(TextRange.create(parameter.getTextLength()-PyNames.OBJECT.length(), parameter.getTextLength()), PyNames.OBJECT); - Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - TemplateManager.getInstance(project).startTemplate(editor, template); + annotateParameter(project, editor, parameter, name, defaultParamValue); } - else { //return type - Callable callable = null; - if (resolvedReference instanceof PyTargetExpression) { - final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue(); - if (assignedValue instanceof PyCallExpression) { - callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext()); - } + else { + annotateReturnType(project, editor, file, resolvedReference); + } + } + + private void annotateParameter(Project project, Editor editor, PyParameter parameter, String name, @Nullable PyExpression defaultParamValue) { + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + + String defaultParamText = defaultParamValue == null? null: defaultParamValue.getText(); + final PyNamedParameter namedParameter = elementGenerator.createParameter(name, defaultParamText, PyNames.OBJECT, LanguageLevel.forElement(parameter)); + assert namedParameter != null; + parameter = (PyParameter)parameter.replace(namedParameter); + parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter); + editor.getCaretModel().moveToOffset(parameter.getTextOffset()); + PyAnnotation annotation = namedParameter.getAnnotation(); + PyExpression annotationValue = annotation.getValue(); + + final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter); + int replacementStart = annotation.getStartOffsetInParent() + annotationValue.getStartOffsetInParent(); + builder.replaceRange(TextRange.create(replacementStart, + replacementStart + annotationValue.getTextLength()), PyNames.OBJECT); + Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + TemplateManager.getInstance(project).startTemplate(editor, template); + } + + private void annotateReturnType(Project project, + Editor editor, + PsiFile file, + PsiElement resolvedReference) { + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + Callable callable = null; + if (resolvedReference instanceof PyTargetExpression) { + final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue(); + if (assignedValue instanceof PyCallExpression) { + callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext()); } - else { - int offset = editor.getCaretModel().getOffset(); - PsiElement elementAt = file.findElementAt(offset - 1); + } + else { + int offset = editor.getCaretModel().getOffset(); + PsiElement elementAt = file.findElementAt(offset - 1); - PyCallExpression callExpression = PyUtil.findProblemElement(editor, file, PyCallExpression.class); - if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) - elementAt = file.findElementAt(offset); + PyCallExpression callExpression = PyUtil.findProblemElement(editor, file, PyCallExpression.class); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) + elementAt = file.findElementAt(offset); - if (callExpression != null) { - callable = callExpression.resolveCalleeFunction(PyResolveContext.defaultContext()); - } - else - callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + if (callExpression != null) { + callable = callExpression.resolveCalleeFunction(PyResolveContext.defaultContext()); } - if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) { - final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText(); - String functionText = functionSignature + - " -> object:"; - for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) { - functionText = functionText + "\n\t" + st.getText(); - } - final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(callable), PyFunction.class, - functionText); - callable = (PyFunction)callable.replace(function); - callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable); + else + callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + } + if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) { + final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText(); + String functionText = functionSignature + + " -> object:"; + for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) { + functionText = functionText + "\n\t" + st.getText(); + } + final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(callable), PyFunction.class, + functionText); + callable = (PyFunction)callable.replace(function); + callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable); - final PyExpression value = ((PyFunction)callable).getAnnotation().getValue(); - final int offset = value.getTextOffset(); + final PyExpression value = ((PyFunction)callable).getAnnotation().getValue(); + final int offset = value.getTextOffset(); - final TemplateBuilder builder = TemplateBuilderFactory.getInstance(). - createTemplateBuilder(value); - builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT); - Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - OpenFileDescriptor descriptor = new OpenFileDescriptor( - project, - value.getContainingFile().getVirtualFile(), - offset - ); - Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); - if (targetEditor != null) { - targetEditor.getCaretModel().moveToOffset(offset); - TemplateManager.getInstance(project).startTemplate(targetEditor, template); - } + final TemplateBuilder builder = TemplateBuilderFactory.getInstance(). + createTemplateBuilder(value); + builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT); + Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + OpenFileDescriptor descriptor = new OpenFileDescriptor( + project, + value.getContainingFile().getVirtualFile(), + offset + ); + Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + if (targetEditor != null) { + targetEditor.getCaretModel().moveToOffset(offset); + TemplateManager.getInstance(project).startTemplate(targetEditor, template); } } } diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index fedab30bf739..9535bded79a4 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -1120,5 +1120,17 @@ public class PyUtil { return null; } + @Nullable + public static PsiElement findNonWhitespaceAtOffset(PsiFile psiFile, int caretOffset) { + PsiElement element = psiFile.findElementAt(caretOffset); + if (element == null) { + return null; + } + while (caretOffset > 0 && element instanceof PsiWhiteSpace) { + caretOffset--; + element = psiFile.findElementAt(caretOffset); + } + return element; + } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java index f7209cf0cb9b..29452169e214 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java @@ -236,7 +236,18 @@ public class PyElementGeneratorImpl extends PyElementGenerator { static int[] PATH_PARAMETER = {0, 3, 1}; public PyNamedParameter createParameter(@NotNull String name) { - return createFromText(LanguageLevel.getDefault(), PyNamedParameter.class, "def f(" + name + "): pass", PATH_PARAMETER); + return createParameter(name, null, null, LanguageLevel.getDefault()); + } + + public PyNamedParameter createParameter(@NotNull String name, @Nullable String defaultValue, @Nullable String annotation, + @NotNull LanguageLevel languageLevel) { + String parameterText = name; + if (annotation != null) + parameterText += " : " + annotation; + if (defaultValue != null) + parameterText += " = " + defaultValue; + + return createFromText(languageLevel, PyNamedParameter.class, "def f(" + parameterText + "): pass", PATH_PARAMETER); } @Override diff --git a/python/testData/intentions/afterTypeAnnotation3.py b/python/testData/intentions/afterTypeAnnotation3.py new file mode 100644 index 000000000000..d0c49d419746 --- /dev/null +++ b/python/testData/intentions/afterTypeAnnotation3.py @@ -0,0 +1,2 @@ +def g(x : object=None ): + return x diff --git a/python/testData/intentions/beforeTypeAnnotation3.py b/python/testData/intentions/beforeTypeAnnotation3.py new file mode 100644 index 000000000000..5d21fd52c050 --- /dev/null +++ b/python/testData/intentions/beforeTypeAnnotation3.py @@ -0,0 +1,2 @@ +def g(x=None): + return x diff --git a/python/testSrc/com/jetbrains/python/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/PyIntentionTest.java index 68e26e3b5590..3bd8b25d6b88 100644 --- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java @@ -254,14 +254,22 @@ public class PyIntentionTest extends PyTestCase { } public void testTypeAssertion() { - doTest(PyBundle.message("INTN.insert.assertion")); + doTestTypeAssertion(); } public void testTypeAssertion1() { //PY-7089 - doTest(PyBundle.message("INTN.insert.assertion")); + doTestTypeAssertion(); } public void testTypeAssertion2() { + doTestTypeAssertion(); + } + + public void testTypeAnnotation3() { //PY-7087 + doTest(PyBundle.message("INTN.specify.type.in.annotation"), LanguageLevel.PYTHON32); + } + + private void doTestTypeAssertion() { doTest(PyBundle.message("INTN.insert.assertion")); }