From 5e89bfdc76d0e4bf674307f628e65f92865bf020 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Tue, 24 Sep 2019 14:28:55 +0900
Subject: [PATCH] Fix fisrt parameter type of a validator method (#76)

---
 resources/META-INF/plugin.xml                 |  1 +
 .../koxudaxi/pydantic/PydanticTypeProvider.kt | 19 +++++++++++++------
 testData/completion/classValidatorCls.py      | 13 +++++++++++++
 .../pydantic/PydanticCompletionTest.kt        |  8 ++++++++
 4 files changed, 35 insertions(+), 6 deletions(-)
 create mode 100644 testData/completion/classValidatorCls.py

diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
index aaa11b23..7543804a 100644
--- 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
index 940daa38..80b8ef72 100644
--- 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
new file mode 100644
index 00000000..21e95771
--- /dev/null
+++ 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
index 333e7811..bae29fc3 100644
--- 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(