diff --git a/python/resources/postfixTemplates/PyMainPostfixTemplate/after.py.template b/python/resources/postfixTemplates/PyMainPostfixTemplate/after.py.template new file mode 100644 index 000000000000..cf0f4f019c15 --- /dev/null +++ b/python/resources/postfixTemplates/PyMainPostfixTemplate/after.py.template @@ -0,0 +1,2 @@ +if __name__ == '__main__': + abs(1) \ No newline at end of file diff --git a/python/resources/postfixTemplates/PyMainPostfixTemplate/before.py.template b/python/resources/postfixTemplates/PyMainPostfixTemplate/before.py.template new file mode 100644 index 000000000000..12b4b3ceca31 --- /dev/null +++ b/python/resources/postfixTemplates/PyMainPostfixTemplate/before.py.template @@ -0,0 +1 @@ +abs(1).main \ No newline at end of file diff --git a/python/resources/postfixTemplates/PyMainPostfixTemplate/description.html b/python/resources/postfixTemplates/PyMainPostfixTemplate/description.html new file mode 100644 index 000000000000..825d342bcb5e --- /dev/null +++ b/python/resources/postfixTemplates/PyMainPostfixTemplate/description.html @@ -0,0 +1,5 @@ + + +Surrounds current statement with if __name__ == '__main__' check. + + \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/postfix/PyMainPostfixTemplate.java b/python/src/com/jetbrains/python/codeInsight/postfix/PyMainPostfixTemplate.java new file mode 100644 index 000000000000..d8f1b117ee0b --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/postfix/PyMainPostfixTemplate.java @@ -0,0 +1,67 @@ +/* + * Copyright 2000-2016 JetBrains s.r.o. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.jetbrains.python.codeInsight.postfix; + +import com.intellij.codeInsight.template.postfix.templates.SurroundPostfixTemplateBase; +import com.intellij.lang.surroundWith.Surrounder; +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.codeStyle.CodeStyleManager; +import com.intellij.util.IncorrectOperationException; +import com.jetbrains.python.psi.LanguageLevel; +import com.jetbrains.python.psi.PyElementGenerator; +import com.jetbrains.python.psi.PyIfStatement; +import com.jetbrains.python.psi.PyStatementList; +import com.jetbrains.python.refactoring.surround.surrounders.statements.PyStatementSurrounder; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +public class PyMainPostfixTemplate extends SurroundPostfixTemplateBase { + + public static final String DESCR = "if __name__ == '__main__': expr"; + + protected PyMainPostfixTemplate() { + super("main", DESCR, PyPostfixUtils.PY_PSI_INFO, PyPostfixUtils.currentStatementSelector()); + } + + @NotNull + @Override + protected Surrounder getSurrounder() { + return new PyStatementSurrounder() { + @Nullable + @Override + protected TextRange surroundStatement(@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement[] elements) + throws IncorrectOperationException { + PyIfStatement ifStatement = PyElementGenerator.getInstance(project).createFromText(LanguageLevel.forElement(elements[0]), PyIfStatement.class, "if __name__ == '__main__':\n expr"); + ifStatement = (PyIfStatement)CodeStyleManager.getInstance(project).reformat(ifStatement); + final PsiElement parent = elements[0].getParent(); + ifStatement = (PyIfStatement) parent.addBefore(ifStatement, elements[0]); + final PyStatementList statementList = ifStatement.getIfPart().getStatementList(); + statementList.addRange(elements[0], elements[elements.length - 1]); + statementList.getFirstChild().delete(); + parent.deleteChildRange(elements[0], elements[elements.length - 1]); + return TextRange.from(statementList.getTextRange().getEndOffset(), 0); + } + + @Override + public String getTemplateDescription() { + return DESCR; + } + }; + } +} diff --git a/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixTemplateProvider.java b/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixTemplateProvider.java index 3bc82e3ea58b..a43bb80b1eed 100644 --- a/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixTemplateProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixTemplateProvider.java @@ -35,7 +35,8 @@ public class PyPostfixTemplateProvider implements PostfixTemplateProvider { new PyWhilePostfixTemplate(), new PyIsNonePostfixTemplate(), new PyIsNotNonePostfixTemplate(), - new PyPrintPostfixTemplate()); + new PyPrintPostfixTemplate(), + new PyMainPostfixTemplate()); } @Override diff --git a/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixUtils.java b/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixUtils.java index e954b4073af0..d4d86443c333 100644 --- a/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixUtils.java +++ b/python/src/com/jetbrains/python/codeInsight/postfix/PyPostfixUtils.java @@ -29,6 +29,7 @@ import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; import java.util.ArrayList; +import java.util.Collections; import java.util.List; public class PyPostfixUtils { @@ -92,4 +93,20 @@ public class PyPostfixUtils { } }; } + + public static PostfixTemplateExpressionSelector currentStatementSelector() { + return new PostfixTemplateExpressionSelectorBase(null) { + @Override + protected List getNonFilteredExpressions(@NotNull PsiElement context, @NotNull Document document, int offset) { + PsiElement elementAtCaret = PsiUtilCore.getElementAtOffset(context.getContainingFile(), offset - 1); + while (elementAtCaret != null) { + if (elementAtCaret instanceof PyStatement) { + return Collections.singletonList(elementAtCaret); + } + elementAtCaret = elementAtCaret.getParent(); + } + return Collections.emptyList(); + } + }; + } } diff --git a/python/testData/postfix/main/print.py b/python/testData/postfix/main/print.py new file mode 100644 index 000000000000..8990e35879a7 --- /dev/null +++ b/python/testData/postfix/main/print.py @@ -0,0 +1 @@ +print("I want to be inside main").main \ No newline at end of file diff --git a/python/testData/postfix/main/print_after.py b/python/testData/postfix/main/print_after.py new file mode 100644 index 000000000000..a120cfa8798e --- /dev/null +++ b/python/testData/postfix/main/print_after.py @@ -0,0 +1,2 @@ +if __name__ == '__main__': + print("I want to be inside main") diff --git a/python/testData/postfix/main/severalStatements.py b/python/testData/postfix/main/severalStatements.py new file mode 100644 index 000000000000..2dc32f4f1cc4 --- /dev/null +++ b/python/testData/postfix/main/severalStatements.py @@ -0,0 +1,2 @@ +print("I want to be inside main").main +print("I want to be inside main too") \ No newline at end of file diff --git a/python/testData/postfix/main/severalStatements_after.py b/python/testData/postfix/main/severalStatements_after.py new file mode 100644 index 000000000000..e73e17bfb5e1 --- /dev/null +++ b/python/testData/postfix/main/severalStatements_after.py @@ -0,0 +1,3 @@ +if __name__ == '__main__': + print("I want to be inside main") +print("I want to be inside main too") \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/postfix/PyMainPostfixTemplateTest.java b/python/testSrc/com/jetbrains/python/postfix/PyMainPostfixTemplateTest.java new file mode 100644 index 000000000000..a7d199e5081c --- /dev/null +++ b/python/testSrc/com/jetbrains/python/postfix/PyMainPostfixTemplateTest.java @@ -0,0 +1,32 @@ +/* + * Copyright 2000-2016 JetBrains s.r.o. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.jetbrains.python.postfix; + +public class PyMainPostfixTemplateTest extends PyPostfixTemplateTestCase { + + public void testPrint() { + doTest(); + } + + public void testSeveralStatements() { + doTest(); + } + + @Override + protected String getTestDataDir() { + return "main/"; + } +}