From bba16d35fba056a5bd6a475d697bfc11312278b6 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Wed, 11 Dec 2019 14:35:02 +0900
Subject: [PATCH 1/6] inspect untyped fields

---
 resources/META-INF/plugin.xml                 |  7 ++++-
 src/com/koxudaxi/pydantic/Pydantic.kt         | 28 +++++++++----------
 .../koxudaxi/pydantic/PydanticInspection.kt   | 28 ++++++++++++++++---
 3 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
index 7993cce5..efa0690a 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">
     <id>com.koxudaxi.pydantic</id>
     <name>Pydantic</name>
-    <version>0.0.28</version>
+    <version>0.0.29</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.0.29</h2>
+    <p>Features, BugFixes</p>
+    <ul>
+        <li>Inspect untyped fields [#93] </li>
+    </ul>
     <h2>version 0.0.28</h2>
     <p>Features, BugFixes</p>
     <ul>
diff --git a/src/com/koxudaxi/pydantic/Pydantic.kt b/src/com/koxudaxi/pydantic/Pydantic.kt
index 92cacd22..c08ca356 100644
--- a/src/com/koxudaxi/pydantic/Pydantic.kt
+++ b/src/com/koxudaxi/pydantic/Pydantic.kt
@@ -55,7 +55,7 @@ val CONFIG_TYPES = mapOf(
         "allow_mutation" to Boolean
 )
 
-internal fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression, context: TypeEvalContext): PyClass? {
+fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression, context: TypeEvalContext): PyClass? {
     val callee = pyCallExpression.callee ?: return null
     val pyType = when (val type = context.getType(callee)) {
         is PyClass -> return type
@@ -65,16 +65,16 @@ internal fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression, co
     return getPyClassTypeByPyTypes(pyType).firstOrNull { isPydanticModel(it.pyClass) }?.pyClass
 }
 
-internal fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument, context: TypeEvalContext): PyClass? {
+fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument, context: TypeEvalContext): PyClass? {
     val pyCallExpression = PsiTreeUtil.getParentOfType(pyKeywordArgument, PyCallExpression::class.java) ?: return null
     return getPyClassByPyCallExpression(pyCallExpression, context)
 }
 
-internal fun isPydanticModel(pyClass: PyClass, context: TypeEvalContext? = null): Boolean {
+fun isPydanticModel(pyClass: PyClass, context: TypeEvalContext? = null): Boolean {
     return (isSubClassOfPydanticBaseModel(pyClass, context) || isPydanticDataclass(pyClass)) && !isPydanticBaseModel(pyClass)
 }
 
-internal fun isPydanticBaseModel(pyClass: PyClass): Boolean {
+fun isPydanticBaseModel(pyClass: PyClass): Boolean {
     return pyClass.qualifiedName == BASE_MODEL_Q_NAME
 }
 
@@ -154,13 +154,13 @@ private fun getAliasedFieldName(field: PyTargetExpression, context: TypeEvalCont
 }
 
 
-internal fun getResolveElements(referenceExpression: PyReferenceExpression, context: TypeEvalContext): Array<ResolveResult> {
+fun getResolveElements(referenceExpression: PyReferenceExpression, context: TypeEvalContext): Array<ResolveResult> {
     val resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context)
     return referenceExpression.getReference(resolveContext).multiResolve(false)
 
 }
 
-internal fun getPyClassTypeByPyTypes(pyType: PyType): List<PyClassType> {
+fun getPyClassTypeByPyTypes(pyType: PyType): List<PyClassType> {
     return when (pyType) {
         is PyUnionType ->
             pyType.members
@@ -174,13 +174,13 @@ internal fun getPyClassTypeByPyTypes(pyType: PyType): List<PyClassType> {
 }
 
 
-internal fun isPydanticSchemaByPsiElement(psiElement: PsiElement, context: TypeEvalContext): Boolean {
+fun isPydanticSchemaByPsiElement(psiElement: PsiElement, context: TypeEvalContext): Boolean {
     PsiTreeUtil.getContextOfType(psiElement, PyClass::class.java)
             ?.let { return isPydanticSchema(it, context) }
     return false
 }
 
-internal fun isPydanticFieldByPsiElement(psiElement: PsiElement): Boolean {
+fun isPydanticFieldByPsiElement(psiElement: PsiElement): Boolean {
     when (psiElement) {
         is PyFunction -> return isPydanticField(psiElement)
         else -> PsiTreeUtil.getContextOfType(psiElement, PyFunction::class.java)
@@ -189,7 +189,7 @@ internal fun isPydanticFieldByPsiElement(psiElement: PsiElement): Boolean {
     return false
 }
 
-internal fun getPydanticVersion(project: Project, context: TypeEvalContext): KotlinVersion? {
+fun getPydanticVersion(project: Project, context: TypeEvalContext): KotlinVersion? {
     val module = project.modules.firstOrNull() ?: return null
     val pythonSdk = module.pythonSdk
     val contextAnchor = ModuleBasedContextAnchor(module)
@@ -210,11 +210,11 @@ internal fun getPydanticVersion(project: Project, context: TypeEvalContext): Kot
     })
 }
 
-internal fun isValidFieldName(name: String): Boolean {
+fun isValidFieldName(name: String): Boolean {
     return name.first() != '_'
 }
 
-internal fun getConfigValue(name: String, value: Any?, context: TypeEvalContext): Any? {
+fun getConfigValue(name: String, value: Any?, context: TypeEvalContext): Any? {
     if (value is PyReferenceExpression) {
         val resolveResults = getResolveElements(value, context)
         val targetExpression = PyUtil.filterTopPriorityResults(resolveResults).firstOrNull() ?: return null
@@ -232,7 +232,7 @@ internal fun getConfigValue(name: String, value: Any?, context: TypeEvalContext)
     return null
 }
 
-internal fun getConfig(pyClass: PyClass, context: TypeEvalContext, setDefault: Boolean): HashMap<String, Any?> {
+fun getConfig(pyClass: PyClass, context: TypeEvalContext, setDefault: Boolean): HashMap<String, Any?> {
     val config = hashMapOf<String, Any?>()
     pyClass.getAncestorClasses(context)
             .reversed()
@@ -265,7 +265,7 @@ internal fun getConfig(pyClass: PyClass, context: TypeEvalContext, setDefault: B
     return config
 }
 
-internal fun getFieldName(field: PyTargetExpression,
+fun getFieldName(field: PyTargetExpression,
                           context: TypeEvalContext,
                           config: HashMap<String, Any?>,
                           pydanticVersion: KotlinVersion?): String? {
@@ -286,7 +286,7 @@ internal fun getFieldName(field: PyTargetExpression,
 }
 
 
-internal fun getPydanticBaseConfig(project: Project, context: TypeEvalContext): PyClass? {
+fun getPydanticBaseConfig(project: Project, context: TypeEvalContext): PyClass? {
     val module = project.modules.firstOrNull() ?: return null
     val pythonSdk = module.pythonSdk
     val contextAnchor = ModuleBasedContextAnchor(module)
diff --git a/src/com/koxudaxi/pydantic/PydanticInspection.kt b/src/com/koxudaxi/pydantic/PydanticInspection.kt
index 620cc7fd..5111ec68 100644
--- a/src/com/koxudaxi/pydantic/PydanticInspection.kt
+++ b/src/com/koxudaxi/pydantic/PydanticInspection.kt
@@ -3,6 +3,7 @@ package com.koxudaxi.pydantic
 import com.intellij.codeInspection.LocalInspectionToolSession
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.codeInspection.ui.MultipleCheckboxOptionsPanel
 import com.intellij.psi.PsiElementVisitor
 import com.jetbrains.python.PyBundle
 import com.jetbrains.python.PyNames
@@ -16,14 +17,17 @@ import com.jetbrains.python.psi.impl.PyTargetExpressionImpl
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.types.PyClassType
 import com.jetbrains.python.psi.types.PyClassTypeImpl
+import javax.swing.JComponent
 
 class PydanticInspection : PyInspection() {
+    var warnUntypedFields = false
+
 
     override fun buildVisitor(holder: ProblemsHolder,
                               isOnTheFly: Boolean,
                               session: LocalInspectionToolSession): PsiElementVisitor = Visitor(holder, session)
 
-    private class Visitor(holder: ProblemsHolder, session: LocalInspectionToolSession) : PyInspectionVisitor(holder, session) {
+    inner class Visitor(holder: ProblemsHolder, session: LocalInspectionToolSession) : PyInspectionVisitor(holder, session) {
 
         override fun visitPyFunction(node: PyFunction?) {
             super.visitPyFunction(node)
@@ -49,7 +53,6 @@ class PydanticInspection : PyInspection() {
             super.visitPyCallExpression(node)
 
             if (node == null) return
-
             inspectPydanticModelCallableExpression(node)
             inspectFromOrm(node)
 
@@ -59,9 +62,10 @@ class PydanticInspection : PyInspection() {
             super.visitPyAssignmentStatement(node)
 
             if (node == null) return
-
+            if (this@PydanticInspection.warnUntypedFields) {
+                inspectWarnUntypedFields(node)
+            }
             inspectReadOnlyProperty(node)
-
         }
 
         private fun inspectPydanticModelCallableExpression(pyCallExpression: PyCallExpression) {
@@ -110,5 +114,21 @@ class PydanticInspection : PyInspection() {
                     ProblemHighlightType.GENERIC_ERROR)
 
         }
+
+        private fun inspectWarnUntypedFields(node: PyAssignmentStatement){
+            val pyClass = node.parent?.parent as? PyClass ?: return
+            if (!isPydanticModel(pyClass, myTypeEvalContext)) return
+            if (node.annotation != null) return
+
+            registerProblem(node,
+                    "Untyped fields disallowed", ProblemHighlightType.WARNING)
+
+        }
+    }
+
+    override fun createOptionsPanel(): JComponent? {
+        val panel = MultipleCheckboxOptionsPanel(this)
+        panel.addCheckbox( "Warning untyped fields", "warnUntypedFields")
+        return panel
     }
 }
\ No newline at end of file

From 2e645856672ecdac691f7482bcaa8458e79da64c Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Thu, 12 Dec 2019 03:08:26 +0900
Subject: [PATCH 2/6] add unittest add description for inspections

---
 .../inspectionDescriptions/PydanticInspection.html  |  5 +++++
 src/com/koxudaxi/pydantic/PydanticInspection.kt     |  7 ++++---
 testData/inspection/warnUntypedFields.py            |  6 ++++++
 testData/inspection/warnUntypedFieldsDisable.py     |  6 ++++++
 .../com/koxudaxi/pydantic/PydanticInspectionTest.kt | 13 +++++++++++++
 5 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 resources/inspectionDescriptions/PydanticInspection.html
 create mode 100644 testData/inspection/warnUntypedFields.py
 create mode 100644 testData/inspection/warnUntypedFieldsDisable.py

diff --git a/resources/inspectionDescriptions/PydanticInspection.html b/resources/inspectionDescriptions/PydanticInspection.html
new file mode 100644
index 00000000..65c6a08d
--- /dev/null
+++ b/resources/inspectionDescriptions/PydanticInspection.html
@@ -0,0 +1,5 @@
+<html>
+<body>
+This inspection checks Pydantic models.
+</body>
+</html>
\ No newline at end of file
diff --git a/src/com/koxudaxi/pydantic/PydanticInspection.kt b/src/com/koxudaxi/pydantic/PydanticInspection.kt
index 5111ec68..a8c6d9fc 100644
--- a/src/com/koxudaxi/pydantic/PydanticInspection.kt
+++ b/src/com/koxudaxi/pydantic/PydanticInspection.kt
@@ -19,16 +19,17 @@ import com.jetbrains.python.psi.types.PyClassType
 import com.jetbrains.python.psi.types.PyClassTypeImpl
 import javax.swing.JComponent
 
-class PydanticInspection : PyInspection() {
-    var warnUntypedFields = false
+var defaultWarnUntypedFields = false
 
+class PydanticInspection : PyInspection() {
+    var warnUntypedFields = defaultWarnUntypedFields
 
     override fun buildVisitor(holder: ProblemsHolder,
                               isOnTheFly: Boolean,
                               session: LocalInspectionToolSession): PsiElementVisitor = Visitor(holder, session)
 
     inner class Visitor(holder: ProblemsHolder, session: LocalInspectionToolSession) : PyInspectionVisitor(holder, session) {
-
+        
         override fun visitPyFunction(node: PyFunction?) {
             super.visitPyFunction(node)
 
diff --git a/testData/inspection/warnUntypedFields.py b/testData/inspection/warnUntypedFields.py
new file mode 100644
index 00000000..6f6cfcaa
--- /dev/null
+++ b/testData/inspection/warnUntypedFields.py
@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    <warning descr="Untyped fields disallowed">a = '123'</warning>
+
diff --git a/testData/inspection/warnUntypedFieldsDisable.py b/testData/inspection/warnUntypedFieldsDisable.py
new file mode 100644
index 00000000..8e166468
--- /dev/null
+++ b/testData/inspection/warnUntypedFieldsDisable.py
@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    a = '123'
+
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticInspectionTest.kt b/testSrc/com/koxudaxi/pydantic/PydanticInspectionTest.kt
index f02de282..f8862ed0 100644
--- a/testSrc/com/koxudaxi/pydantic/PydanticInspectionTest.kt
+++ b/testSrc/com/koxudaxi/pydantic/PydanticInspectionTest.kt
@@ -49,4 +49,17 @@ open class PydanticInspectionTest : PydanticInspectionBase() {
     fun testReadOnlyProperty() {
         doTest()
     }
+
+    fun testWarnUntypedFieldsDisable() {
+        doTest()
+    }
+
+    fun testWarnUntypedFields() {
+        try {
+            defaultWarnUntypedFields = true
+            doTest()
+        } finally {
+            defaultWarnUntypedFields = false
+        }
+    }
 }

From 1695452faafad87b147a4dae1a6cfce276d1dbf4 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Thu, 12 Dec 2019 03:32:23 +0900
Subject: [PATCH 3/6] add testcase

---
 testData/inspection/warnUntypedFields.py        | 2 ++
 testData/inspection/warnUntypedFieldsDisable.py | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/testData/inspection/warnUntypedFields.py b/testData/inspection/warnUntypedFields.py
index 6f6cfcaa..f0078098 100644
--- a/testData/inspection/warnUntypedFields.py
+++ b/testData/inspection/warnUntypedFields.py
@@ -4,3 +4,5 @@
 class A(BaseModel):
     <warning descr="Untyped fields disallowed">a = '123'</warning>
 
+class B(BaseModel):
+    b = '123'
diff --git a/testData/inspection/warnUntypedFieldsDisable.py b/testData/inspection/warnUntypedFieldsDisable.py
index 8e166468..e98aa5cf 100644
--- a/testData/inspection/warnUntypedFieldsDisable.py
+++ b/testData/inspection/warnUntypedFieldsDisable.py
@@ -4,3 +4,6 @@
 class A(BaseModel):
     a = '123'
 
+
+class B(BaseModel):
+    b = '123'

From 3b72ad7537406d3e9b50410f31ff9c78e437b585 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Thu, 12 Dec 2019 03:51:36 +0900
Subject: [PATCH 4/6] add testcase

---
 testData/inspection/warnUntypedFields.py        | 2 +-
 testData/inspection/warnUntypedFieldsDisable.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/testData/inspection/warnUntypedFields.py b/testData/inspection/warnUntypedFields.py
index f0078098..1c68b62e 100644
--- a/testData/inspection/warnUntypedFields.py
+++ b/testData/inspection/warnUntypedFields.py
@@ -5,4 +5,4 @@ class A(BaseModel):
     <warning descr="Untyped fields disallowed">a = '123'</warning>
 
 class B(BaseModel):
-    b = '123'
+    b: str = '123'
diff --git a/testData/inspection/warnUntypedFieldsDisable.py b/testData/inspection/warnUntypedFieldsDisable.py
index e98aa5cf..270f99f2 100644
--- a/testData/inspection/warnUntypedFieldsDisable.py
+++ b/testData/inspection/warnUntypedFieldsDisable.py
@@ -6,4 +6,4 @@ class A(BaseModel):
 
 
 class B(BaseModel):
-    b = '123'
+    b: str = '123'

From b8a3eb3c7016c40b5d8470faa34ada257bd27467 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Thu, 12 Dec 2019 13:53:03 +0900
Subject: [PATCH 5/6] add testcase

---
 testData/inspection/warnUntypedFields.py        | 8 ++++++++
 testData/inspection/warnUntypedFieldsDisable.py | 8 ++++++++
 2 files changed, 16 insertions(+)

diff --git a/testData/inspection/warnUntypedFields.py b/testData/inspection/warnUntypedFields.py
index 1c68b62e..86c16dd7 100644
--- a/testData/inspection/warnUntypedFields.py
+++ b/testData/inspection/warnUntypedFields.py
@@ -4,5 +4,13 @@
 class A(BaseModel):
     <warning descr="Untyped fields disallowed">a = '123'</warning>
 
+
 class B(BaseModel):
     b: str = '123'
+
+
+class C:
+    c = '123'
+
+class D:
+    d
\ No newline at end of file
diff --git a/testData/inspection/warnUntypedFieldsDisable.py b/testData/inspection/warnUntypedFieldsDisable.py
index 270f99f2..5db2cc8f 100644
--- a/testData/inspection/warnUntypedFieldsDisable.py
+++ b/testData/inspection/warnUntypedFieldsDisable.py
@@ -7,3 +7,11 @@ class A(BaseModel):
 
 class B(BaseModel):
     b: str = '123'
+
+
+class C:
+    c = '123'
+
+
+class D:
+    d
\ No newline at end of file

From 5066c35e882deed9ecb05f57fb692716067c3d7b Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Thu, 12 Dec 2019 23:09:16 +0900
Subject: [PATCH 6/6] add getPyClassByAttribute

---
 src/com/koxudaxi/pydantic/Pydantic.kt                      | 4 ++++
 src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt | 6 +++---
 src/com/koxudaxi/pydantic/PydanticInspection.kt            | 7 ++++---
 src/com/koxudaxi/pydantic/PydanticTypeProvider.kt          | 2 +-
 testData/inspection/warnUntypedFields.py                   | 7 ++++++-
 testData/inspection/warnUntypedFieldsDisable.py            | 7 ++++++-
 6 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/com/koxudaxi/pydantic/Pydantic.kt b/src/com/koxudaxi/pydantic/Pydantic.kt
index c08ca356..1c928de8 100644
--- a/src/com/koxudaxi/pydantic/Pydantic.kt
+++ b/src/com/koxudaxi/pydantic/Pydantic.kt
@@ -291,4 +291,8 @@ fun getPydanticBaseConfig(project: Project, context: TypeEvalContext): PyClass?
     val pythonSdk = module.pythonSdk
     val contextAnchor = ModuleBasedContextAnchor(module)
     return BASE_CONFIG_QUALIFIED_NAME.resolveToElement(QNameResolveContext(contextAnchor, pythonSdk, context)) as? PyClass
+}
+
+fun getPyClassByAttribute(pyPsiElement: PsiElement?): PyClass? {
+    return pyPsiElement?.parent?.parent as? PyClass
 }
\ No newline at end of file
diff --git a/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt b/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt
index 954be2a8..db1f63b9 100644
--- a/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt
+++ b/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt
@@ -240,9 +240,9 @@ class PydanticCompletionContributor : CompletionContributor() {
 
 
         override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
-            val configClass = parameters.position.parent?.parent?.parent?.parent as? PyClass ?: return
+            val configClass = getPyClassByAttribute(parameters.position.parent?.parent) ?: return
             if (!isConfigClass(configClass)) return
-            val pydanticModel = configClass.parent?.parent as? PyClass ?:return
+            val pydanticModel = getPyClassByAttribute(configClass) ?:return
             if (!isPydanticModel(pydanticModel)) return
             val typeEvalContext = parameters.getTypeEvalContext()
 
@@ -264,7 +264,7 @@ class PydanticCompletionContributor : CompletionContributor() {
         override val icon: Icon = AllIcons.Nodes.Class
 
         override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
-            val pydanticModel = parameters.position.parent?.parent?.parent?.parent as? PyClass ?: return
+            val pydanticModel = getPyClassByAttribute(parameters.position.parent?.parent) ?: return
             if (!isPydanticModel(pydanticModel)) return
             if (pydanticModel.findNestedClass("Config", false) != null) return
             val element = PrioritizedLookupElement.withGrouping(
diff --git a/src/com/koxudaxi/pydantic/PydanticInspection.kt b/src/com/koxudaxi/pydantic/PydanticInspection.kt
index a8c6d9fc..ab461e55 100644
--- a/src/com/koxudaxi/pydantic/PydanticInspection.kt
+++ b/src/com/koxudaxi/pydantic/PydanticInspection.kt
@@ -29,11 +29,12 @@ class PydanticInspection : PyInspection() {
                               session: LocalInspectionToolSession): PsiElementVisitor = Visitor(holder, session)
 
     inner class Visitor(holder: ProblemsHolder, session: LocalInspectionToolSession) : PyInspectionVisitor(holder, session) {
-        
+
         override fun visitPyFunction(node: PyFunction?) {
             super.visitPyFunction(node)
 
-            val pyClass = node?.parent?.parent as? PyClass ?: return
+            if (node == null) return
+            val pyClass = getPyClassByAttribute(node) ?: return
             if (!isPydanticModel(pyClass, myTypeEvalContext) || !isValidatorMethod(node)) return
             val paramList = node.parameterList
             val params = paramList.parameters
@@ -117,7 +118,7 @@ class PydanticInspection : PyInspection() {
         }
 
         private fun inspectWarnUntypedFields(node: PyAssignmentStatement){
-            val pyClass = node.parent?.parent as? PyClass ?: return
+            val pyClass = getPyClassByAttribute(node) ?: return
             if (!isPydanticModel(pyClass, myTypeEvalContext)) return
             if (node.annotation != null) return
 
diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
index 189a1038..aea5ac3d 100644
--- a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
+++ b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
@@ -19,7 +19,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
 
     override fun getReferenceType(referenceTarget: PsiElement, context: TypeEvalContext, anchor: PsiElement?): Ref<PyType>? {
         if (referenceTarget is PyTargetExpression) {
-            val pyClass = referenceTarget.parent?.parent?.parent as? PyClass ?: return null
+            val pyClass = getPyClassByAttribute(referenceTarget.parent) ?: return null
             if (!isPydanticModel(pyClass, context)) return null
             val name = referenceTarget.name ?: return null
             getRefTypeFromFieldName(name, context, pyClass)?.let { return it }
diff --git a/testData/inspection/warnUntypedFields.py b/testData/inspection/warnUntypedFields.py
index 86c16dd7..1dd01104 100644
--- a/testData/inspection/warnUntypedFields.py
+++ b/testData/inspection/warnUntypedFields.py
@@ -13,4 +13,9 @@ class C:
     c = '123'
 
 class D:
-    d
\ No newline at end of file
+    d
+
+def e():
+    ee = '123'
+
+f = '123'
diff --git a/testData/inspection/warnUntypedFieldsDisable.py b/testData/inspection/warnUntypedFieldsDisable.py
index 5db2cc8f..c6e3c11e 100644
--- a/testData/inspection/warnUntypedFieldsDisable.py
+++ b/testData/inspection/warnUntypedFieldsDisable.py
@@ -14,4 +14,9 @@ class C:
 
 
 class D:
-    d
\ No newline at end of file
+    d
+
+def e():
+    ee = '123'
+
+f = '123'