PsiElement (@NotNull PsiNamedElement element) { if (element instanceof PyClass || element instanceof PyFunction) { return element; } else if (element instanceof PyTargetExpression && element.getParent() instanceof PyAssignmentStatement) { return element.getParent(); } return null; }