From 1878dbb1e19fdeb1842db5e9b6a6d65ad5170d9e Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Wed, 6 Jan 2021 19:05:58 +0900
Subject: [PATCH] Fix custom root inspection

---
 resources/META-INF/plugin.xml                   | 7 ++++++-
 src/com/koxudaxi/pydantic/PydanticInspection.kt | 5 +++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
index a8ae3b57..30d887ff 100644
--- a/resources/META-INF/plugin.xml
+++ b/resources/META-INF/plugin.xml
@@ -1,9 +1,14 @@
 <idea-plugin url="https://github.com/koxudaxi/pydantic-pycharm-plugin" require-restart="true">
     <id>com.koxudaxi.pydantic</id>
     <name>Pydantic</name>
-    <version>0.1.18</version>
+    <version>0.1.19</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.1.19</h2>
+    <p>BugFixes</p>
+    <ul>
+        <li>Fix custom root inspection [#232]</li>
+    </ul>
     <h2>version 0.1.18</h2>
     <p>Features</p>
     <ul>
diff --git a/src/com/koxudaxi/pydantic/PydanticInspection.kt b/src/com/koxudaxi/pydantic/PydanticInspection.kt
index bad1771e..c25a4b2a 100644
--- a/src/com/koxudaxi/pydantic/PydanticInspection.kt
+++ b/src/com/koxudaxi/pydantic/PydanticInspection.kt
@@ -124,10 +124,11 @@ class PydanticInspection : PyInspection() {
 
         private fun inspectCustomRootField(node: PyAssignmentStatement) {
             val pyClass = getPyClassByAttribute(node) ?: return
-            if (!isPydanticModel(pyClass, true, myTypeEvalContext)) return
+            if (!isPydanticModel(pyClass, false, myTypeEvalContext)) return
             val fieldName = (node.leftHandSideExpression as? PyTargetExpressionImpl)?.text ?: return
             if (fieldName.startsWith('_')) return
-            pyClass.findClassAttribute("__root__", true, myTypeEvalContext) ?: return
+            val rootModel = pyClass.findClassAttribute("__root__", true, myTypeEvalContext)?.containingClass ?: return
+            if (!isPydanticModel(rootModel, false, myTypeEvalContext)) return
             registerProblem(node,
                 "__root__ cannot be mixed with other fields", ProblemHighlightType.WARNING)
         }