PY-19606 main postfix template

2026-09-27 10:03:11 +07:00 · 2016-06-01 02:10:10 -07:00
parent c5ce062f81
commit a0c901b890
11 changed files with 134 additions and 1 deletions
@@ -0,0 +1,2 @@
+if __name__ == '__main__':
+    abs(1)
@@ -0,0 +1 @@
+<spot>abs(1)</spot>.main
@@ -0,0 +1,5 @@
+<html>
+<body>
+Surrounds current statement with if __name__ == '__main__' check.
+</body>
+</html>
@@ -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;
+      }
+    };
+  }
+}
@@ -35,7 +35,8 @@ public class PyPostfixTemplateProvider implements PostfixTemplateProvider {
                                                     new PyWhilePostfixTemplate(),
                                                     new PyIsNonePostfixTemplate(),
                                                     new PyIsNotNonePostfixTemplate(),
-                                                     new PyPrintPostfixTemplate());
+                                                     new PyPrintPostfixTemplate(),
+                                                     new PyMainPostfixTemplate());
  }

  @Override
@@ -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<PsiElement> 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();
+      }
+    };
+  }
 }
@@ -0,0 +1 @@
+print("I want to be inside main").main<caret>
@@ -0,0 +1,2 @@
+if __name__ == '__main__':
+    print("I want to be inside main")
@@ -0,0 +1,2 @@
+print("I want to be inside main").main<caret>
+print("I want to be inside main too")
@@ -0,0 +1,3 @@
+if __name__ == '__main__':
+    print("I want to be inside main")
+print("I want to be inside main too")
@@ -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/";
+  }
+}
				`@@ -0,0 +1 @@`
				`print("I want to be inside main").main<caret>`