PsiElement (@NotNull DataContext dataContext) { Project project = CommonDataKeys.PROJECT.getData(dataContext); if (project == null) return null; PsiElement element = CommonDataKeys.PSI_ELEMENT.getData(dataContext); if (element == null) { Editor editor = CommonDataKeys.EDITOR.getData(dataContext); if (editor != null) { PsiFile file = PsiDocumentManager.getInstance(project).getPsiFile(editor.getDocument()); if (file == null) return null; element = TargetElementUtil.findTargetElement(editor, TargetElementUtil.ELEMENT_NAME_ACCEPTED | TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED | TargetElementUtil.LOOKUP_ITEM_ACCEPTED); if (element instanceof PyFunction || element instanceof PyClass || element instanceof PyFile) { return element; } element = file.findElementAt(editor.getCaretModel().getOffset()); } } return PsiTreeUtil.getNonStrictParentOfType(element, PyFunction.class, PyClass.class, PyFile.class); }