Support ClassVar

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.1.12</version>
+    <version>0.1.13</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.1.13</h2>
+    <p>Features</p>
+    <ul>
+        <li>Support ClassVar [#188]</li>
+    </ul>
     <h2>version 0.1.12</h2>
     <p>Features</p>
     <ul>

src/com/koxudaxi/pydantic/Pydantic.kt

@@ -13,10 +13,7 @@ import com.jetbrains.extensions.QNameResolveContext
 import com.jetbrains.extensions.resolveToElement
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
 import com.jetbrains.python.psi.*
-import com.jetbrains.python.psi.impl.PyCallExpressionImpl
-import com.jetbrains.python.psi.impl.PyReferenceExpressionImpl
-import com.jetbrains.python.psi.impl.PyStarArgumentImpl
-import com.jetbrains.python.psi.impl.PyTargetExpressionImpl
+import com.jetbrains.python.psi.impl.*
 import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.resolve.PyResolveUtil
 import com.jetbrains.python.psi.types.*
@@ -251,8 +248,17 @@ fun getPydanticVersion(project: Project, context: TypeEvalContext): KotlinVersio
     })
 }
 
-fun isValidFieldName(name: String): Boolean {
-    return !name.startsWith('_')
+fun isValidField(field: PyTargetExpression): Boolean {
+    if (!isValidFieldName(field.name)) return false
+
+    val annotation = field.annotation?.value ?: return true
+    val qualifier = (annotation as? PySubscriptionExpression)?.qualifier ?: return true
+    // TODO Support a variable.
+    return qualifier.text != "ClassVar"
+}
+
+fun isValidFieldName(name: String?): Boolean {
+    return name?.startsWith('_') == false
 }
 
 fun getConfigValue(name: String, value: Any?, context: TypeEvalContext): Any? {
@@ -354,7 +360,7 @@ fun createPyClassTypeImpl(qualifiedName: String, project: Project, context: Type
 
 fun getPydanticPyClass(pyCallExpression: PyCallExpression, context: TypeEvalContext): PyClass? {
     val pyClass = getPyClassByPyCallExpression(pyCallExpression, false, context) ?: return null
-    if(!isPydanticModel(pyClass, false, context)) return null
+    if (!isPydanticModel(pyClass, false, context)) return null
     if ((pyCallExpression.callee as? PyReferenceExpressionImpl)?.isQualified == true) return null
     return pyClass
 }
@@ -366,9 +372,10 @@ fun getParentOfPydanticCallableExpression(file: PsiFile, offset: Int, context: T
     }
     return pyCallExpression
 }
+
 fun getPydanticCallExpressionAtCaret(file: PsiFile, editor: Editor, context: TypeEvalContext): PyCallExpression? {
-    return getParentOfPydanticCallableExpression(file, editor.caretModel.offset, context) ?:
-            getParentOfPydanticCallableExpression(file, editor.caretModel.offset - 1, context)
+    return getParentOfPydanticCallableExpression(file, editor.caretModel.offset, context)
+            ?: getParentOfPydanticCallableExpression(file, editor.caretModel.offset - 1, context)
 }
 
 

src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt

@@ -90,7 +90,7 @@ class PydanticCompletionContributor : CompletionContributor() {
             val pydanticVersion = getPydanticVersion(pyClass.project, typeEvalContext)
             getClassVariables(pyClass, typeEvalContext)
                     .filter { it.name != null }
-                    .filter { isValidFieldName(it.name!!) }
+                    .filter { isValidField(it) }
                     .filter { !isDataclass || isInInit(it) }
                     .forEach {
                         val elementName = getLookupNameFromFieldName(it, typeEvalContext, pydanticVersion, config)
@@ -141,20 +141,16 @@ class PydanticCompletionContributor : CompletionContributor() {
                     .filter { isPydanticModel(it, true) }
                     .forEach {
                         fieldElements.addAll(it.classAttributes
-                                .filter { attribute
-                                    ->
-                                    attribute.name?.let { name -> isValidFieldName(name) } ?: false
+                                .filter { attribute ->
+                                    isValidField(attribute)
                                 }
                                 .mapNotNull { attribute -> attribute?.name })
                     }
 
 
 
             fieldElements.addAll(pyClass.classAttributes
-                    .filter { attribute
-                        ->
-                        attribute.name?.let { name -> isValidFieldName(name) } ?: false
-                    }
+                    .filter { isValidField(it) }
                     .mapNotNull { attribute -> attribute?.name })
 
             result.runRemainingContributors(parameters)

src/com/koxudaxi/pydantic/PydanticTypeProvider.kt

@@ -251,7 +251,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
                                   config: HashMap<String, Any?>,
                                   typed: Boolean = true,
                                   isDataclass: Boolean = false): PyCallableParameter? {
-        if (field.name == null || !isValidFieldName(field.name!!)) return null
+        if (!isValidField(field)) return null
         if (!hasAnnotationValue(field) && !field.hasAssignedValue()) return null // skip fields that are invalid syntax
 
         val defaultValueFromField = getDefaultValueForParameter(field, ellipsis, context, pydanticVersion, isDataclass)

testData/completion/fieldIgnore.py

@@ -1,14 +1,16 @@
 from builtins import *
+from typing import ClassVar
 from pydantic import BaseModel
 
 
 class A(BaseModel):
     _abc: str = str('abc')
     __cde: str = str('abc')
-
+    efg: ClassVar[str] = str('abc')
 
 class B(A):
     _efg: str = str('abc')
     __hij: str = str('abc')
+    klm: ClassVar[str] = str('abc')
 
 A().<caret>

testData/mock/stub/typing/__init__.py

@@ -23,3 +23,8 @@ class List:
     @classmethod
     def __getitem__(cls, item):
         pass
+
+class ClassVar:
+    @classmethod
+    def __getitem__(cls, item):
+        pass

testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt

@@ -439,6 +439,7 @@ open class PydanticCompletionTest : PydanticTestCase() {
     fun testFieldIgnore() {
         doFieldTest(
                 listOf(
+                        Pair("efg", "A"),
                         Pair("___slots__", "BaseModel")
                 )
         )