From d8a85181df8a021c1d4334d7445e5e31b4aa36f5 Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Fri, 26 Oct 2012 14:59:54 +0400
Subject: [PATCH] fixed  PY-7860 Doctest: respect python language level in
 doctest syntax highlighting

---
 .../src/com/jetbrains/python/psi/LanguageLevel.java   |  6 +++++-
 .../doctest/PyDocstringParserDefinition.java          | 11 +++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/LanguageLevel.java b/python/psi-api/src/com/jetbrains/python/psi/LanguageLevel.java
index 04590189fce0..c2073f4434d5 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/LanguageLevel.java
+++ b/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;
diff --git a/python/src/com/jetbrains/python/documentation/doctest/PyDocstringParserDefinition.java b/python/src/com/jetbrains/python/documentation/doctest/PyDocstringParserDefinition.java
index 86448fee9cbb..420701f45095 100644
--- a/python/src/com/jetbrains/python/documentation/doctest/PyDocstringParserDefinition.java
+++ b/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();
+      }
     };
   }
 }