From 8c0fe707e529cf5fc43173126b137804873e660f Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Thu, 24 Nov 2016 14:15:02 +0300 Subject: [PATCH] PY-21204 Create a new class in host file instead of injected fragment Additionally, I removed usages of deprecated API in CreateClassQuickFix and utilize SmartPointer to hold the reference to PSI element. --- .../quickfix/CreateClassQuickFix.java | 48 ++++++++++++++----- .../inspections/AddClassFromTypeComment.py | 2 + .../AddClassFromTypeComment_after.py | 6 +++ .../com/jetbrains/python/PyQuickFixTest.java | 5 ++ 4 files changed, 48 insertions(+), 13 deletions(-) create mode 100644 python/testData/inspections/AddClassFromTypeComment.py create mode 100644 python/testData/inspections/AddClassFromTypeComment_after.py diff --git a/python/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java index 2eb31cdd83f3..b5a6002e7591 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/CreateClassQuickFix.java @@ -20,10 +20,20 @@ import com.intellij.codeInsight.template.TemplateBuilder; import com.intellij.codeInsight.template.TemplateBuilderFactory; import com.intellij.codeInspection.LocalQuickFix; import com.intellij.codeInspection.ProblemDescriptor; +import com.intellij.lang.injection.InjectedLanguageManager; +import com.intellij.openapi.fileEditor.FileEditor; +import com.intellij.openapi.fileEditor.FileEditorManager; +import com.intellij.openapi.fileEditor.TextEditor; import com.intellij.openapi.project.Project; import com.intellij.psi.PsiElement; +import com.intellij.psi.PsiLanguageInjectionHost; +import com.intellij.psi.SmartPointerManager; +import com.intellij.psi.SmartPsiElementPointer; import com.jetbrains.python.PyNames; -import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.LanguageLevel; +import com.jetbrains.python.psi.PyClass; +import com.jetbrains.python.psi.PyElementGenerator; +import com.jetbrains.python.psi.PyFile; import com.jetbrains.python.psi.impl.PyPsiUtils; import org.jetbrains.annotations.NotNull; @@ -32,11 +42,16 @@ import org.jetbrains.annotations.NotNull; */ public class CreateClassQuickFix implements LocalQuickFix { private final String myClassName; - private final PsiElement myAnchor; + private final SmartPsiElementPointer myAnchor; public CreateClassQuickFix(String className, PsiElement anchor) { myClassName = className; - myAnchor = anchor; + + final Project project = anchor.getProject(); + final PsiLanguageInjectionHost injectionHost = InjectedLanguageManager.getInstance(project).getInjectionHost(anchor); + final PsiElement notInjectedAnchor = injectionHost != null ? injectionHost : anchor; + + myAnchor = SmartPointerManager.getInstance(project).createSmartPsiElementPointer(notInjectedAnchor); } @NotNull @@ -49,22 +64,22 @@ public class CreateClassQuickFix implements LocalQuickFix { @NotNull public String getFamilyName() { - return "Create Class"; + return "Create class"; } public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) { - PsiElement anchor = myAnchor; - PyPsiUtils.assertValid(anchor); - if (!anchor.isValid()) { + PsiElement anchor = myAnchor.getElement(); + if (anchor == null || !anchor.isValid()) { return; } + if (!(anchor instanceof PyFile)) { - while(!(anchor.getParent() instanceof PyFile)) { - anchor = anchor.getParent(); - } + anchor = PyPsiUtils.getParentRightBefore(anchor, anchor.getContainingFile()); + assert anchor != null; } + PyClass pyClass = PyElementGenerator.getInstance(project).createFromText(LanguageLevel.getDefault(), PyClass.class, - "class " + myClassName + "(object):\n pass"); + "class " + myClassName + "(object):\n pass"); if (anchor instanceof PyFile) { pyClass = (PyClass) anchor.add(pyClass); } @@ -72,10 +87,17 @@ public class CreateClassQuickFix implements LocalQuickFix { pyClass = (PyClass) anchor.getParent().addBefore(pyClass, anchor); } pyClass = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(pyClass); - TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(pyClass); + final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(pyClass); builder.replaceElement(pyClass.getSuperClassExpressions() [0], "object"); builder.replaceElement(pyClass.getStatementList(), PyNames.PASS); - builder.run(); + + + final FileEditor editor = FileEditorManager.getInstance(project).getSelectedEditor(anchor.getContainingFile().getVirtualFile()); + if (!(editor instanceof TextEditor)) { + return; + } + + builder.run(((TextEditor)editor).getEditor(), false); } } diff --git a/python/testData/inspections/AddClassFromTypeComment.py b/python/testData/inspections/AddClassFromTypeComment.py new file mode 100644 index 000000000000..185a91eae58d --- /dev/null +++ b/python/testData/inspections/AddClassFromTypeComment.py @@ -0,0 +1,2 @@ +def f(): + x = None # type: MyClass diff --git a/python/testData/inspections/AddClassFromTypeComment_after.py b/python/testData/inspections/AddClassFromTypeComment_after.py new file mode 100644 index 000000000000..e687a07c2059 --- /dev/null +++ b/python/testData/inspections/AddClassFromTypeComment_after.py @@ -0,0 +1,6 @@ +class MyClass(object): + pass + + +def f(): + x = None # type: MyClass diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java index 7c09d8c24b7d..6227cf0ade1a 100644 --- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java @@ -202,6 +202,11 @@ public class PyQuickFixTest extends PyTestCase { doInspectionTest("AddClass.py", PyUnresolvedReferencesInspection.class, "Create class 'Xyzzy'", true, true); } + // PY-21204 + public void testAddClassFromTypeComment() { + doInspectionTest(PyUnresolvedReferencesInspection.class, "Create class 'MyClass'", true, true); + } + // PY-1602 public void testAddFunctionToModule() { doInspectionTest(