Merge pull request #271 from koxudaxi/improve_supporting_dynamic_model

Improve supporting dynamic model
koxudaxi · Apr 12, 2021 · 77ac9e1 · 77ac9e1
2 parents 42b666d + 38cba16
commit 77ac9e1
Show file tree

Hide file tree

Showing 20 changed files with 587 additions and 199 deletions.
diff --git a/src/com/koxudaxi/pydantic/Pydantic.kt b/src/com/koxudaxi/pydantic/Pydantic.kt
@@ -281,8 +281,8 @@ fun getPydanticVersion(project: Project, context: TypeEvalContext): KotlinVersio
     val version = getPsiElementByQualifiedName(VERSION_QUALIFIED_NAME, project, context) as? PyTargetExpression
             ?: return null
     val versionString = (version.findAssignedValue()?.lastChild?.firstChild?.nextSibling as? PyStringLiteralExpression)?.stringValue
-            ?: return null
-    return pydanticVersionCache.getOrPut(versionString, {
+        ?:(version.findAssignedValue() as? PyStringLiteralExpressionImpl)?.stringValue ?: return null
+    return pydanticVersionCache.getOrPut(versionString) {
         val versionList = versionString.split(VERSION_SPLIT_PATTERN).map { it.toIntOrNull() ?: 0 }
         val pydanticVersion = when {
             versionList.size == 1 -> KotlinVersion(versionList[0], 0)
@@ -292,7 +292,7 @@ fun getPydanticVersion(project: Project, context: TypeEvalContext): KotlinVersio
         } ?: KotlinVersion(0, 0)
         pydanticVersionCache[versionString] = pydanticVersion
         pydanticVersion
-    })
+    }
 }
 
 fun isValidField(field: PyTargetExpression, context: TypeEvalContext): Boolean {

diff --git a/src/com/koxudaxi/pydantic/PydanticDynamicModel.kt b/src/com/koxudaxi/pydantic/PydanticDynamicModel.kt
@@ -1,15 +1,39 @@
 package com.koxudaxi.pydantic
 
+import com.intellij.icons.AllIcons
 import com.intellij.lang.ASTNode
+import com.jetbrains.python.codeInsight.PyCustomMember
 import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyElement
+import com.jetbrains.python.psi.PyExpression
 import com.jetbrains.python.psi.impl.PyClassImpl
+import com.jetbrains.python.psi.types.PyCallableParameter
 import com.jetbrains.python.psi.types.PyClassLikeType
 import com.jetbrains.python.psi.types.TypeEvalContext
 
-class PydanticDynamicModel(astNode: ASTNode, val baseModel: PyClass) : PyClassImpl(astNode) {
+class PydanticDynamicModel(astNode: ASTNode, val baseModel: PyClass, val attributes: Map<String, Attribute>) : PyClassImpl(astNode) {
+    val members: List<PyCustomMember> =  attributes.values.map { it.pyCustomMember }
+    private val memberResolver: Map<String, PyElement> = attributes.entries.filterNot { it.value.isInAncestor } .associate { it.key to it.value.pyElement }
+
+    fun resolveMember(name: String): PyElement? = memberResolver[name]
+
     override fun getSuperClassTypes(context: TypeEvalContext): MutableList<PyClassLikeType> {
         return baseModel.getType(context)?.let {
             mutableListOf(it)
         } ?: mutableListOf()
     }
+    data class Attribute(val pyCallableParameter: PyCallableParameter, val pyCustomMember: PyCustomMember, val pyElement: PyElement, val isInAncestor: Boolean)
+
+    companion object {
+        fun createAttribute(name: String, parameter: PyCallableParameter, originalPyExpression: PyExpression, context: TypeEvalContext, isInAncestor: Boolean): Attribute {
+            val type = parameter.getType(context)
+            return Attribute(parameter,
+                    PyCustomMember(name, null) { type }
+                    .toPsiElement(originalPyExpression)
+                    .withIcon(AllIcons.Nodes.Field),
+                    originalPyExpression,
+                    isInAncestor
+            )
+        }
+    }
 }
diff --git a/src/com/koxudaxi/pydantic/PydanticDynamicModelClassType.kt b/src/com/koxudaxi/pydantic/PydanticDynamicModelClassType.kt
@@ -1,10 +1,15 @@
 package com.koxudaxi.pydantic
 
-
-import com.jetbrains.python.codeInsight.PyCustomMember
-import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.types.PyCallableType
+import com.jetbrains.python.psi.types.PyCallableTypeImpl
 import com.jetbrains.python.psi.types.PyClassTypeImpl
 
-class PydanticDynamicModelClassType(source: PyClass, isDefinition: Boolean, val members: List<PyCustomMember>, private val memberResolver: Map<String, PyElement>) : PyClassTypeImpl(source, isDefinition) {
-    fun resolveMember(name: String): PyElement? = memberResolver[name]
+class PydanticDynamicModelClassType(private val source: PydanticDynamicModel, isDefinition: Boolean) :
+    PyClassTypeImpl(source, isDefinition) {
+
+    val pyCallableType: PyCallableType
+        get() = PyCallableTypeImpl(
+            source.attributes.values.map { attribute -> attribute.pyCallableParameter },
+            this.toInstance()
+        )
 }
diff --git a/src/com/koxudaxi/pydantic/PydanticDynamicModelMemberProvider.kt b/src/com/koxudaxi/pydantic/PydanticDynamicModelMemberProvider.kt
@@ -7,14 +7,19 @@ import com.jetbrains.python.psi.types.*
 
 class PydanticDynamicModelMemberProvider : PyClassMembersProviderBase() {
     override fun resolveMember(type: PyClassType, name: String, location: PsiElement?, resolveContext: PyResolveContext): PsiElement? {
-        if (type is PydanticDynamicModelClassType) {
-            type.resolveMember(name)?.let { return it }
-        }
+        val pyClass =  type.pyClass
+        if (pyClass is PydanticDynamicModel && !type.isDefinition)
+            pyClass.resolveMember(name)?.let { return it }
         return super.resolveMember(type, name, location, resolveContext)
     }
 
     override fun getMembers(clazz: PyClassType?, location: PsiElement?, context: TypeEvalContext): MutableCollection<PyCustomMember> {
-        if (clazz !is PydanticDynamicModelClassType) return mutableListOf()
-        return clazz.members.toMutableList()
+        if (clazz == null || clazz.isDefinition) return mutableListOf()
+        val pyClass =  clazz.pyClass
+        return if (pyClass is PydanticDynamicModel) {
+            pyClass.members.toMutableList()
+        } else {
+            mutableListOf()
+        }
     }
 }