Merge pull request #227 from koxudaxi/support_custom_root

Support custom root field

resources/META-INF/plugin.xml

@@ -1,9 +1,14 @@
 <idea-plugin url="https://github.com/koxudaxi/pydantic-pycharm-plugin" require-restart="true">
     <id>com.koxudaxi.pydantic</id>
     <name>Pydantic</name>
-    <version>0.1.17</version>
+    <version>0.1.18</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.1.18</h2>
+    <p>Features</p>
+    <ul>
+        <li>Support custom root field [#227]</li>
+    </ul>
     <h2>version 0.1.17</h2>
     <p>Features</p>
     <ul>

src/com/koxudaxi/pydantic/Pydantic.kt

@@ -101,6 +101,8 @@ val CONFIG_TYPES = mapOf(
         "keep_untouched" to ConfigType.LIST_PYTYPE
 )
 
+const val CUSTOM_ROOT_FIELD = "__root__"
+
 fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression, includeDataclass: Boolean, context: TypeEvalContext): PyClass? {
     val callee = pyCallExpression.callee ?: return null
     val pyType = when (val type = context.getType(callee)) {
@@ -290,7 +292,7 @@ fun isValidField(field: PyTargetExpression): Boolean {
 }
 
 fun isValidFieldName(name: String?): Boolean {
-    return name?.startsWith('_') == false
+    return name?.let { !it.startsWith('_') || it == CUSTOM_ROOT_FIELD } ?: false
 }
 
 fun getConfigValue(name: String, value: Any?, context: TypeEvalContext): Any? {

src/com/koxudaxi/pydantic/PydanticInspection.kt

@@ -63,6 +63,7 @@ class PydanticInspection : PyInspection() {
             if (pydanticConfigService.currentWarnUntypedFields) {
                 inspectWarnUntypedFields(node)
             }
+            inspectCustomRootField(node)
             inspectReadOnlyProperty(node)
         }
 
@@ -115,11 +116,21 @@ class PydanticInspection : PyInspection() {
             val pyClass = getPyClassByAttribute(node) ?: return
             if (!isPydanticModel(pyClass, true, myTypeEvalContext)) return
             if (node.annotation != null) return
-            if ((node.leftHandSideExpression as? PyTargetExpressionImpl)?.text?.startsWith("_") == true) return
+            if (!isValidFieldName((node.leftHandSideExpression as? PyTargetExpressionImpl)?.text)) return
             registerProblem(node,
                     "Untyped fields disallowed", ProblemHighlightType.WARNING)
 
         }
+
+        private fun inspectCustomRootField(node: PyAssignmentStatement) {
+            val pyClass = getPyClassByAttribute(node) ?: return
+            if (!isPydanticModel(pyClass, true, myTypeEvalContext)) return
+            val fieldName = (node.leftHandSideExpression as? PyTargetExpressionImpl)?.text ?: return
+            if (fieldName.startsWith('_')) return
+            pyClass.findClassAttribute("__root__", true, myTypeEvalContext) ?: return
+            registerProblem(node,
+                "__root__ cannot be mixed with other fields", ProblemHighlightType.WARNING)
+        }
     }
 
 //    override fun createOptionsPanel(): JComponent? {

src/com/koxudaxi/pydantic/PydanticTypeProvider.kt

@@ -200,7 +200,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
         }
         argumentWithoutModelName
                 .filter { it is PyKeywordArgument || (it as? PyStarArgumentImpl)?.isKeyword == true }
-                .filterNot { it.name?.startsWith("_") == true || it.name == "model_name" }
+                .filter { isValidFieldName(it.name) || it.name != "model_name" }
                 .forEach {
                     val parameter = fieldToParameter(it, context, hashMapOf(), typed)!!
                     parameter.name?.let { name ->

testData/completion/fieldCustomRoot.py

@@ -0,0 +1,7 @@
+from builtins import *
+from pydantic import BaseModel
+
+class A(BaseModel):
+    __root__: str
+
+A().<caret>

testData/completion/keywordArgumentCustomRoot.py

@@ -0,0 +1,12 @@
+from builtins import *
+
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    __root__: str
+
+class B(A):
+    hij: str
+
+A(<caret>)

testData/inspection/customRoot.py

@@ -0,0 +1,29 @@
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    __root__ = 'xyz'
+
+
+class B(BaseModel):
+    a = 'xyz'
+
+
+class C(BaseModel):
+    __root__ = 'xyz'
+    <warning descr="__root__ cannot be mixed with other fields">b = 'xyz'</warning>
+
+
+class D(BaseModel):
+    __root__ = 'xyz'
+    _c = 'xyz'
+    __c = 'xyz'
+
+class E:
+    __root__ = 'xyz'
+    e = 'xyz'
+
+def f():
+    __root__ = 'xyz'
+    g = 'xyz'
+

testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt

@@ -33,6 +33,14 @@ open class PydanticCompletionTest : PydanticTestCase() {
         )
     }
 
+    fun testKeywordArgumentCustomRoot() {
+        doFieldTest(
+                listOf(
+                        Pair("__root__=", "str A")
+                )
+        )
+    }
+
     fun testKeywordArgumentDot() {
         doFieldTest(
                 listOf(Pair("___slots__", "BaseModel"))
@@ -445,6 +453,14 @@ open class PydanticCompletionTest : PydanticTestCase() {
                 )
         )
     }
+    fun testFieldCustomRoot() {
+        doFieldTest(
+                listOf(
+                        Pair("__root__", "str A"),
+                        Pair("___slots__", "BaseModel")
+                )
+        )
+    }
 
     fun testFieldOptional() {
         doFieldTest(

testSrc/com/koxudaxi/pydantic/PydanticInspectionTest.kt

@@ -59,4 +59,8 @@ open class PydanticInspectionTest : PydanticInspectionBase() {
         pydanticConfigService.warnUntypedFields = true
         doTest()
     }
+
+    fun testCustomRoot() {
+        doTest()
+    }
 }