PsiElement (@NotNull final PsiElement element) { assertValid(element); final PsiFile file = element.getContainingFile(); if (file instanceof PyExpressionCodeFragment) { final PsiElement context = file.getContext(); return context != null ? context : element; } else { return element; } }