Improve field detections

koxudaxi · Apr 14, 2024 · 323b6dd · 323b6dd
1 parent fafd533
commit 323b6dd
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 23 deletions.
diff --git a/src/com/koxudaxi/pydantic/PydanticInspection.kt b/src/com/koxudaxi/pydantic/PydanticInspection.kt
@@ -3,7 +3,6 @@ package com.koxudaxi.pydantic
 import com.intellij.codeInspection.LocalInspectionToolSession
 import com.intellij.codeInspection.ProblemHighlightType
 import com.intellij.codeInspection.ProblemsHolder
-import com.intellij.diff.comparison.trimEnd
 import com.intellij.psi.PsiElementVisitor
 import com.jetbrains.python.PyNames
 import com.jetbrains.python.codeInsight.stdlib.PyDataclassTypeProvider
@@ -16,7 +15,10 @@ import com.jetbrains.python.psi.impl.PyCallExpressionImpl
 import com.jetbrains.python.psi.impl.PyEvaluator
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl
 import com.jetbrains.python.psi.resolve.PyResolveContext
-import com.jetbrains.python.psi.types.*
+import com.jetbrains.python.psi.types.PyCallableType
+import com.jetbrains.python.psi.types.PyClassType
+import com.jetbrains.python.psi.types.PyTypeChecker
+import com.jetbrains.python.psi.types.TypeEvalContext
 
 
 class PydanticInspection : PyInspection() {
@@ -426,21 +428,33 @@ class PydanticInspection : PyInspection() {
         private fun inspectModelAttribute(node: PyQualifiedExpression) {
             val qualifier = node.qualifier ?: return
             val name = node.name ?: return
-            val pyClass = (myTypeEvalContext.getType(qualifier) as? PyClassType)?.pyClass ?: return
+            val pyClassType = myTypeEvalContext.getType(qualifier) as? PyClassType ?: return
+            val pyClass = pyClassType.pyClass
             if (!isPydanticModel(pyClass, false, myTypeEvalContext)) return
+            if (pyClass.findNestedClass(name, true) is PyClass) return
             if (pyClass.findProperty(name, true, myTypeEvalContext) != null) return
             if (pyClass.findMethodByName(name, true, myTypeEvalContext) != null) return
+            val field = pyClass.findClassAttribute(name, true, myTypeEvalContext)
+            if (field is PyAnnotationOwner && PyTypingTypeProvider.isClassVar(field, myTypeEvalContext)) return
+
             val pydanticVersion = PydanticCacheService.getVersion(pyClass.project)
+
+            // Check private field or model fields
+            if (name.startsWith("_") || pydanticVersion.isV2 && name.startsWith(MODEL_FIELD_PREFIX)) return
+
             val config = getConfig(pyClass, myTypeEvalContext, true)
-            getAncestorPydanticModels(pyClass, true, myTypeEvalContext).forEach {
-                if (hasAttribute(it, config, pydanticVersion.isV2, name)) return
+            if (pyClassType.isDefinition) {
+                if(field == null && node.reference?.resolve() is PyTargetExpression) return
+            } else {
+                getAncestorPydanticModels(pyClass, true, myTypeEvalContext).forEach {
+                    if (hasAttribute(it, config, pydanticVersion.isV2, name)) return
+                }
+                if (hasAttribute(pyClass, config, pydanticVersion.isV2, name)) return
             }
-            if (hasAttribute(pyClass, config, pydanticVersion.isV2, name)) return
             registerProblem(node.node.lastChildNode.psi, "Unresolved attribute reference '${name}' for class '${pyClass.name}' ")
         }
         private fun hasAttribute(pyClass: PyClass, config: HashMap<String, Any?>, isV2: Boolean, name: String): Boolean =
-            getPydanticField(pyClass, myTypeEvalContext, config, isV2, false, name)
-                    .any() || pyClass.findNestedClass(name, true) is PyClass
+            getPydanticField(pyClass, myTypeEvalContext, config, isV2, false, name).any()
 
     }
 

diff --git a/testData/inspection/readOnlyProperty.py b/testData/inspection/readOnlyProperty.py
@@ -45,9 +45,9 @@ class G(BaseModel):
     class Config:
         allow_mutation=False
 G.abc =<EOLError descr="Expression expected"></EOLError>
-G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()
-<error descr="Cannot assign to function call">G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()</error> = 'efg'
-
+G.abc = "test"
+G.abc.lower()
+<error descr="Cannot assign to function call">G.abc.lower()</error> = 'efg'
 
 class H:
     class Config:

diff --git a/testData/inspectionv18/readOnlyProperty.py b/testData/inspectionv18/readOnlyProperty.py
@@ -45,9 +45,9 @@ class G(BaseModel):
     class Config:
         allow_mutation=False
 G.abc =<EOLError descr="Expression expected"></EOLError>
-G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()
-<error descr="Cannot assign to function call">G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()</error> = 'efg'
-
+G.abc = "test"
+G.abc.lower()
+<error descr="Cannot assign to function call">G.abc.lower()</error> = 'efg'
 
 class H:
     class Config:

diff --git a/testData/inspectionv18/readOnlyPropertyFrozen.py b/testData/inspectionv18/readOnlyPropertyFrozen.py
@@ -45,9 +45,9 @@ class G(BaseModel):
     class Config:
         frozen=False
 G.abc =<EOLError descr="Expression expected"></EOLError>
-G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()
-<error descr="Cannot assign to function call">G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()</error> = 'efg'
-
+G.abc = "test"
+G.abc.lower()
+<error descr="Cannot assign to function call">G.abc.lower()</error> = 'efg'
 
 class H:
     class Config:

diff --git a/testData/inspectionv2/modelAttribute.py b/testData/inspectionv2/modelAttribute.py
@@ -8,13 +8,14 @@ class B(BaseModel):
     class Inner(BaseModel):
         user: str
     a: int = 1
-#     _url: str = "https://someurl"
+    _url: str = "https://someurl"
     TEST: ClassVar[str] = "Hello World"
     def f(self):
+        self._url
         self.Inner
         self.<warning descr="Unresolved attribute reference 'fake' for class 'B'">fake</warning>
         self.a
-#         self._url
+B.<warning descr="Unresolved attribute reference 'a' for class 'B'">a</warning>
 b = B(a=1)
 b.<warning descr="Unresolved attribute reference 'fake' for class 'B'">fake</warning>
 B.Inner

diff --git a/testData/inspectionv2/readOnlyPropertyFrozenConfig.py b/testData/inspectionv2/readOnlyPropertyFrozenConfig.py
@@ -45,8 +45,9 @@ class G(BaseModel):
     class Config:
         frozen=False
 G.abc =<EOLError descr="Expression expected"></EOLError>
-G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()
-<error descr="Cannot assign to function call">G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()</error> = 'efg'
+G.abc = "test"
+G.abc.lower()
+<error descr="Cannot assign to function call">G.abc.lower()</error> = 'efg'
 
 
 class H:

diff --git a/testData/inspectionv2/readOnlyPropertyFrozenConfigDict.py b/testData/inspectionv2/readOnlyPropertyFrozenConfigDict.py
@@ -40,8 +40,9 @@ class F(D):
 class G(BaseModel):
     model_config = ConfigDict(frozen=False)
 G.abc =<EOLError descr="Expression expected"></EOLError>
-G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()
-<error descr="Cannot assign to function call">G.<warning descr="Unresolved attribute reference 'abc' for class 'G'">abc</warning>.lower()</error> = 'efg'
+G.abc = "test"
+G.abc.lower()
+<error descr="Cannot assign to function call">G.abc.lower()</error> = 'efg'
 
 
 class H: