Fix False positive for non-decorator use of validator (#818)

koxudaxi · Oct 12, 2023 · 87c92a9 · 87c92a9
1 parent 25cae75
commit 87c92a9
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 5 deletions.
diff --git a/src/com/koxudaxi/pydantic/PydanticReferenceContributor.kt b/src/com/koxudaxi/pydantic/PydanticReferenceContributor.kt
@@ -16,7 +16,10 @@ class PydanticReferenceContributor : PsiReferenceContributor() {
                     if (element !is PyStringLiteralExpression) {
                         return false
                     }
-                    return isValidatorField(element, TypeEvalContext.userInitiated(element.project, element.containingFile))
+                    return isValidatorField(
+                        element,
+                        TypeEvalContext.userInitiated(element.project, element.containingFile)
+                    )
                 }
             }
     }
@@ -41,16 +44,29 @@ class PydanticReferenceContributor : PsiReferenceContributor() {
             provider: PsiReferenceProvider
         ) {
             registrar.registerReferenceProvider(
-                PlatformPatterns.psiElement(
-                    PyStringLiteralExpression::class.java
-                ).withParents(
+                PlatformPatterns.or(
+                    PlatformPatterns.psiElement(
+                        PyStringLiteralExpression::class.java
+                    ).withParents(
                         PyArgumentList::class.java,
                         PyCallExpression::class.java,
                         PyDecorator::class.java,
                         PyDecoratorList::class.java,
                         PyFunction::class.java,
                         PyStatementList::class.java,
-                        PyClass::class.java).with(Holder.VALIDATOR_FIELD), provider
+                        PyClass::class.java
+                    ).with(Holder.VALIDATOR_FIELD),
+                    PlatformPatterns.psiElement(
+                        PyStringLiteralExpression::class.java
+                    ).withParents(
+                        PyArgumentList::class.java,
+                        PyCallExpression::class.java,
+                        PyCallExpression::class.java,
+                        PyAssignmentStatement::class.java,
+                        PyStatementList::class.java,
+                        PyClass::class.java
+                    ).with(Holder.VALIDATOR_FIELD)
+                ), provider
             )
         }
     }

diff --git a/testData/inspectionv18/validatorField.py b/testData/inspectionv18/validatorField.py
@@ -0,0 +1,21 @@
+from pydantic import BaseModel, validator
+
+
+class A(BaseModel):
+    foo: int
+    bar: float
+
+    _validate_stuff = validator(
+        "bar",
+        allow_reuse=True,
+    )(do_stuff)
+
+
+class B(BaseModel):
+    foo: int
+    bar: float
+
+    _validate_stuff = validator(
+        <error descr="Cannot find field 'abc'">"abc"</error>,
+        allow_reuse=True,
+    )(do_stuff)
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticInspectionV18Test.kt b/testSrc/com/koxudaxi/pydantic/PydanticInspectionV18Test.kt
@@ -28,4 +28,8 @@ open class PydanticInspectionV18Test : PydanticInspectionBase(version = "v18") {
     fun testDefaultFactory() {
         doTest()
     }
+
+    fun testValidatorField() {
+        doTest()
+    }
 }