@NotNull
  public static PsiElement getRealContext(@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;
    }
  }
@NotNull
  public static PsiElement getRealContext(@NotNull final PsiElement element)