diff --git a/python/src/com/jetbrains/python/inspections/PyClassHasNoInitInspection.java b/python/src/com/jetbrains/python/inspections/PyClassHasNoInitInspection.java index 92ef945255b6..89e8e95cf26b 100644 --- a/python/src/com/jetbrains/python/inspections/PyClassHasNoInitInspection.java +++ b/python/src/com/jetbrains/python/inspections/PyClassHasNoInitInspection.java @@ -27,7 +27,6 @@ import com.jetbrains.python.psi.PyClass; import com.jetbrains.python.psi.PyFunction; import com.jetbrains.python.psi.types.PyClassLikeType; import com.jetbrains.python.psi.types.PyClassType; -import com.jetbrains.python.psi.types.PyClassTypeImpl; import org.jetbrains.annotations.Nls; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -77,7 +76,7 @@ public class PyClassHasNoInitInspection extends PyInspection { final PyFunction init = node.findInitOrNew(true); if (init == null) { registerProblem(node.getNameIdentifier(), PyBundle.message("INSP.class.has.no.init"), - new AddMethodQuickFix("__init__", new PyClassTypeImpl(node, false), false)); + new AddMethodQuickFix("__init__", node.getName(), false)); } } } diff --git a/python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java b/python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java index 730cc804ca9c..eee1e9cf3b14 100644 --- a/python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java +++ b/python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java @@ -751,7 +751,7 @@ public class PyUnresolvedReferencesInspection extends PyInspection { PyClass cls = ((PyClassType)type).getPyClass(); if (!PyBuiltinCache.getInstance(element).isBuiltin(cls)) { if (element.getParent() instanceof PyCallExpression) { - actions.add(new AddMethodQuickFix(refText, (PyClassType)type, true)); + actions.add(new AddMethodQuickFix(refText, cls.getName(), true)); } else if (!(reference instanceof PyOperatorReference)) { actions.add(new AddFieldQuickFix(refText, "None", type.getName())); diff --git a/python/src/com/jetbrains/python/inspections/quickfix/AddMethodQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/AddMethodQuickFix.java index 30d99103b887..a138555d3693 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/AddMethodQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/AddMethodQuickFix.java @@ -21,9 +21,15 @@ 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.openapi.editor.Editor; +import com.intellij.openapi.fileEditor.FileEditorManager; +import com.intellij.openapi.fileEditor.OpenFileDescriptor; import com.intellij.openapi.project.Project; import com.intellij.openapi.ui.MessageType; +import com.intellij.openapi.vfs.VirtualFile; 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; @@ -31,6 +37,7 @@ import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.ParamHelper; import com.jetbrains.python.psi.impl.PyFunctionBuilder; import com.jetbrains.python.psi.types.PyClassType; +import com.jetbrains.python.psi.types.PyClassTypeImpl; import com.jetbrains.python.psi.types.PyType; import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; @@ -44,20 +51,20 @@ import static com.jetbrains.python.psi.PyUtil.sure; */ public class AddMethodQuickFix implements LocalQuickFix { - private PyClassType myQualifierType; + private final String myClassName; private final boolean myReplaceUsage; private String myIdentifier; - public AddMethodQuickFix(String identifier, PyClassType qualifierType, + public AddMethodQuickFix(String identifier, String className, boolean replaceUsage) { myIdentifier = identifier; - myQualifierType = qualifierType; + myClassName = className; myReplaceUsage = replaceUsage; } @NotNull public String getName() { - return PyBundle.message("QFIX.NAME.add.method.$0.to.class.$1", myIdentifier, myQualifierType.getName()); + return PyBundle.message("QFIX.NAME.add.method.$0.to.class.$1", myIdentifier, myClassName); } @NotNull @@ -68,46 +75,46 @@ public class AddMethodQuickFix implements LocalQuickFix { public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) { try { // there can be no name clash, else the name would have resolved, and it hasn't. - PsiElement problem_elt = descriptor.getPsiElement(); - PyClass cls = myQualifierType.getPyClass(); - boolean call_by_class = myQualifierType.isDefinition(); - String item_name = myIdentifier; - sure(item_name); - PyStatementList cls_stmt_list = cls.getStatementList(); - sure(FileModificationService.getInstance().preparePsiElementForWrite(cls_stmt_list)); + final PsiElement problemElement = descriptor.getPsiElement(); + final PyClassType type = getClassType(problemElement); + if (type == null) return; + final PyClass cls = type.getPyClass(); + boolean callByClass = type.isDefinition(); + PyStatementList clsStmtList = cls.getStatementList(); + sure(FileModificationService.getInstance().preparePsiElementForWrite(clsStmtList)); // try to at least match parameter count // TODO: get parameter style from code style - PyFunctionBuilder builder = new PyFunctionBuilder(item_name); - PsiElement pe = problem_elt.getParent(); - String deco_name = null; // set to non-null to add a decorator + PyFunctionBuilder builder = new PyFunctionBuilder(myIdentifier); + PsiElement pe = problemElement.getParent(); + String decoratorName = null; // set to non-null to add a decorator PyExpression[] args = new PyExpression[0]; if (pe instanceof PyCallExpression) { PyArgumentList arglist = ((PyCallExpression)pe).getArgumentList(); - sure(arglist); + if (arglist == null) return; args = arglist.getArguments(); } - boolean made_instance = false; - if (call_by_class) { + boolean madeInstance = false; + if (callByClass) { if (args.length > 0) { - PyType first_arg_type = TypeEvalContext.userInitiated(cls.getContainingFile()).getType(args[0]); - if (first_arg_type instanceof PyClassType && ((PyClassType)first_arg_type).getPyClass().isSubclass(cls)) { + PyType firstArgType = TypeEvalContext.userInitiated(cls.getContainingFile()).getType(args[0]); + if (firstArgType instanceof PyClassType && ((PyClassType)firstArgType).getPyClass().isSubclass(cls)) { // class, first arg ok: instance method builder.parameter("self"); // NOTE: might use a name other than 'self', according to code style. - made_instance = true; + madeInstance = true; } } - if (! made_instance) { // class, first arg absent or of different type: classmethod + if (!madeInstance) { // class, first arg absent or of different type: classmethod builder.parameter("cls"); // NOTE: might use a name other than 'cls', according to code style. - deco_name = PyNames.CLASSMETHOD; + decoratorName = PyNames.CLASSMETHOD; } } else { // instance method builder.parameter("self"); // NOTE: might use a name other than 'self', according to code style. } - boolean skip_first = call_by_class && made_instance; // ClassFoo.meth(foo_instance) + boolean skipFirst = callByClass && madeInstance; // ClassFoo.meth(foo_instance) for (PyExpression arg : args) { - if (skip_first) { - skip_first = false; + if (skipFirst) { + skipFirst = false; continue; } if (arg instanceof PyKeywordArgument) { // foo(bar) -> def foo(self, bar_1) @@ -121,16 +128,18 @@ public class AddMethodQuickFix implements LocalQuickFix { builder.parameter("param"); } } - PyFunction meth = builder.buildFunction(project, LanguageLevel.getDefault()); - if (deco_name != null) { + PyFunction method = builder.buildFunction(project, LanguageLevel.getDefault()); + if (decoratorName != null) { PyElementGenerator generator = PyElementGenerator.getInstance(project); - PyDecoratorList deco_list = generator.createFromText(LanguageLevel.getDefault(), PyDecoratorList.class, "@" + deco_name + "\ndef foo(): pass", new int[]{0, 0}); - meth.addBefore(deco_list, meth.getFirstChild()); // in the very beginning + PyDecoratorList decoratorList = generator + .createFromText(LanguageLevel.getDefault(), PyDecoratorList.class, "@" + decoratorName + "\ndef foo(): pass", new int[]{0, 0}); + method.addBefore(decoratorList, method.getFirstChild()); // in the very beginning } - meth = (PyFunction)PyUtil.addElementToStatementList(meth, cls_stmt_list, PyNames.INIT.equals(meth.getName())); - if (myReplaceUsage) - showTemplateBuilder(meth); + method = (PyFunction)PyUtil.addElementToStatementList(method, clsStmtList, PyNames.INIT.equals(method.getName())); + if (myReplaceUsage) { + showTemplateBuilder(method, problemElement.getContainingFile()); + } } catch (IncorrectOperationException ignored) { // we failed. tell about this @@ -138,7 +147,18 @@ public class AddMethodQuickFix implements LocalQuickFix { } } - private static void showTemplateBuilder(PyFunction method) { + private static PyClassType getClassType(@NotNull final PsiElement problemElement) { + if ((problemElement instanceof PyQualifiedExpression)) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier == null) return null; + final PyType type = TypeEvalContext.userInitiated(problemElement.getContainingFile()).getType(qualifier); + return type instanceof PyClassType ? (PyClassType)type : null; + } + final PyClass pyClass = PsiTreeUtil.getParentOfType(problemElement, PyClass.class); + return pyClass != null ? new PyClassTypeImpl(pyClass, false) : null; + } + + private static void showTemplateBuilder(PyFunction method, PsiFile file) { method = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(method); final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(method); @@ -152,9 +172,13 @@ public class AddMethodQuickFix implements LocalQuickFix { ); final PyStatementList statementList = method.getStatementList(); - if (statementList == null) return; builder.replaceElement(statementList, PyNames.PASS); - builder.run(); + final VirtualFile virtualFile = file.getVirtualFile(); + if (virtualFile == null) return; + final Editor editor = FileEditorManager.getInstance(file.getProject()).openTextEditor( + new OpenFileDescriptor(file.getProject(), virtualFile), true); + if (editor == null) return; + builder.run(editor, false); } }