fixed PY-7860 Doctest: respect python language level in doctest syntax highlighting

python/psi-api/src/com/jetbrains/python/psi/LanguageLevel.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python.psi;
 
+import com.intellij.injected.editor.VirtualFileWindow;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -101,6 +102,9 @@ public enum LanguageLevel {
 
   @NotNull
   public static LanguageLevel forFile(@NotNull VirtualFile virtualFile) {
+    if (virtualFile instanceof VirtualFileWindow)
+      virtualFile = ((VirtualFileWindow)virtualFile).getDelegate();
+
     // Most of the cases should be handled by this one, PyLanguageLevelPusher pushes folders only
     final VirtualFile folder = virtualFile.getParent();
     if (folder != null) {
@@ -111,7 +115,7 @@ public enum LanguageLevel {
       // However this allows us to setup language level per file manually
       // in case when it is LightVirtualFile
       final LanguageLevel level = virtualFile.getUserData(KEY);
-      if (level != null) return level;      
+      if (level != null) return level;
 
       if (ApplicationManager.getApplication().isUnitTestMode()) {
         final LanguageLevel languageLevel = FORCE_LANGUAGE_LEVEL;

python/src/com/jetbrains/python/documentation/doctest/PyDocstringParserDefinition.java

@@ -1,14 +1,17 @@
 package com.jetbrains.python.documentation.doctest;
 
 import com.intellij.lang.PsiParser;
+import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.lexer.Lexer;
 import com.intellij.openapi.fileTypes.FileType;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.FileViewProvider;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.tree.IFileElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.jetbrains.python.PythonParserDefinition;
+import com.jetbrains.python.psi.LanguageLevel;
 import com.jetbrains.python.psi.impl.PyFileImpl;
 import org.jetbrains.annotations.NotNull;
 
@@ -50,6 +53,14 @@ public class PyDocstringParserDefinition extends PythonParserDefinition {
       public FileType getFileType() {
         return PyDocstringFileType.INSTANCE;
       }
+
+      @Override
+      public LanguageLevel getLanguageLevel() {
+        final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(getProject());
+        final PsiLanguageInjectionHost host = languageManager.getInjectionHost(this);
+        if (host != null) return LanguageLevel.forElement(host.getContainingFile());
+        return super.getLanguageLevel();
+      }
     };
   }
 }