Fix fisrt parameter type of a validator method (#76)

koxudaxi · Sep 24, 2019 · 5e89bfd · 5e89bfd
1 parent 62900a4
commit 5e89bfd
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 6 deletions.
diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
@@ -7,6 +7,7 @@
     <h2>version 0.0.22</h2>
     <p>Features, BugFixes</p>
     <ul>
+        <li>Fix first parameter type of a validator method [#76] </li>
         <li>Fix auto-completion for Fields [#75] </li>
         <li>Improve to insert validate methods [#74] </li>
     </ul>

diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
@@ -28,13 +28,20 @@ class PydanticTypeProvider : PyTypeProviderBase() {
     }
 
     override fun getParameterType(param: PyNamedParameter, func: PyFunction, context: TypeEvalContext): Ref<PyType>? {
-        if (!param.isPositionalContainer && !param.isKeywordContainer && param.annotationValue == null && func.name == "__init__") {
-            val pyClass = func.containingClass ?: return null
-            if (!isPydanticModel(pyClass, context)) return null
-            val name = param.name ?: return null
-            getRefTypeFromFieldName(name, context, pyClass)?.let { return it }
+        return when {
+            !param.isPositionalContainer && !param.isKeywordContainer && param.annotationValue == null && func.name == "__init__" -> {
+                val pyClass = func.containingClass ?: return null
+                if (!isPydanticModel(pyClass, context)) return null
+                val name = param.name ?: return null
+                getRefTypeFromFieldName(name, context, pyClass)?.let { it }
+            }
+            param.isSelf && isValidatorMethod(func) -> {
+                val pyClass = func.containingClass ?: return null
+                if (!isPydanticModel(pyClass, context)) return null
+                Ref.create(context.getType(pyClass))
+            }
+            else -> null
         }
-        return null
     }
 
     private fun getRefTypeFromFieldName(name: String, context: TypeEvalContext, pyClass: PyClass): Ref<PyType>? {

diff --git a/testData/completion/classValidatorCls.py b/testData/completion/classValidatorCls.py
@@ -0,0 +1,13 @@
+,from builtins import *
+
+from pydantic import BaseModel, validator
+
+
+class A(BaseModel):
+    abc: str
+
+    @validator('abc')
+    def test(cls):
+        return cls.<caret>
+
+
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt b/testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt
@@ -431,6 +431,14 @@ open class PydanticCompletionTest : PydanticTestCase() {
         )
     }
 
+    fun testClassValidatorCls() {
+        doFieldTest(
+                listOf(
+                        Pair("___slots__", "BaseModel")
+                )
+        )
+    }
+
     fun testMethodSelf() {
         doFieldTest(
                 listOf(