Merge pull request #160 from koxudaxi/fix_inserting_arguments

fix inserting arguments

src/com/koxudaxi/pydantic/Pydantic.kt

@@ -342,13 +342,17 @@ fun getPydanticPyClass(pyCallExpression: PyCallExpression, context: TypeEvalCont
     if ((pyCallExpression.callee as? PyReferenceExpressionImpl)?.isQualified == true) return null
     return pyClass
 }
-fun getTopmostParentOfPyCallExpression(file: PsiFile, offset: Int): PyCallExpression? =
-        PsiTreeUtil.getTopmostParentOfType(file.findElementAt(offset), PyCallExpression::class.java)
 
-fun getPyCallExpressionAtCaret(file: PsiFile, editor: Editor): PyCallExpression? {
-    return getTopmostParentOfPyCallExpression(file, editor.caretModel.offset)
-            ?: getTopmostParentOfPyCallExpression(file, editor.caretModel.offset - 1)
-            ?:return null
+fun getParentOfPydanticCallableExpression(file: PsiFile, offset: Int, context: TypeEvalContext): PyCallExpression? {
+    var pyCallExpression: PyCallExpression? = PsiTreeUtil.getParentOfType(file.findElementAt(offset), PyCallExpression::class.java, true)
+    while (pyCallExpression != null && getPydanticPyClass(pyCallExpression, context) == null) {
+        pyCallExpression = PsiTreeUtil.getParentOfType(pyCallExpression, PyCallExpression::class.java, true)
+    }
+    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)
 }
 
 

src/com/koxudaxi/pydantic/PydanticInsertArgumentsQuickFix.kt

@@ -30,17 +30,17 @@ class PydanticInsertArgumentsQuickFix(private val onlyRequired: Boolean) : Local
     @Throws(IncorrectOperationException::class)
     override fun invoke(project: Project, editor: Editor, file: PsiFile) {
         ApplicationManager.getApplication().runWriteAction {
-            getPyCallExpressionAtCaret(file, editor)?.let { runFix(project, file, it) }
+            val context = TypeEvalContext.userInitiated(project, file)
+            getPydanticCallExpressionAtCaret(file, editor, context)?.let { runFix(project, file, it, context) }
         }
     }
 
     override fun startInWriteAction(): Boolean = true
 
-    fun runFix(project: Project, file: PsiFile, originalElement: PsiElement): PyCallExpression? {
+    fun runFix(project: Project, file: PsiFile, originalElement: PsiElement, context: TypeEvalContext): PyCallExpression? {
         if (originalElement !is PyCallExpression) return null
         if (file !is PyFile) return null
         val newEl = originalElement.copy() as PyCallExpression
-        val context = TypeEvalContext.userInitiated(originalElement.project, originalElement.containingFile)
         val unFilledArguments = getPydanticUnFilledArguments(null, originalElement, pydanticTypeProvider, context).let {
             when {
                 onlyRequired -> it.filter { arguments -> arguments.required }
@@ -61,7 +61,9 @@ class PydanticInsertArgumentsQuickFix(private val onlyRequired: Boolean) : Local
     }
 
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-        runFix(project, descriptor.psiElement.containingFile, descriptor.psiElement)
+        descriptor.psiElement.containingFile.let {
+            runFix(project, it, descriptor.psiElement, TypeEvalContext.userInitiated(project, it))
+        }
     }
 
 }

testData/insertargumentsquickfix/nestedOtherObject.py

@@ -0,0 +1,14 @@
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class A(BaseModel):
+    a: int
+    b: int = ...
+    c: int = 123
+    d: int = Field(123)
+    e: int = Field(...)
+    f: Optional[int]
+
+
+str(A(a=int(<caret>), c=2))

testData/insertargumentsquickfix/nestedOtherObject_after.py

@@ -0,0 +1,14 @@
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class A(BaseModel):
+    a: int
+    b: int = ...
+    c: int = 123
+    d: int = Field(123)
+    e: int = Field(...)
+    f: Optional[int]
+
+
+str(A(a=int(), c=2, b=, d=123, e=, f=))

testData/insertargumentsquickfix/nestedPartArguments.py

@@ -0,0 +1,14 @@
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class A(BaseModel):
+    a: int
+    b: int = ...
+    c: int = 123
+    d: int = Field(123)
+    e: int = Field(...)
+    f: Optional[int]
+
+
+str(A<caret>(a=1, c=2))

testData/insertargumentsquickfix/nestedPartArguments_after.py

@@ -0,0 +1,14 @@
+from typing import Optional
+from pydantic import BaseModel, Field
+
+
+class A(BaseModel):
+    a: int
+    b: int = ...
+    c: int = 123
+    d: int = Field(123)
+    e: int = Field(...)
+    f: Optional[int]
+
+
+str(A(a=1, c=2, b=, d=123, e=, f=))

testSrc/com/koxudaxi/pydantic/PydanticInsertArgumentsQuickFixTest.kt

@@ -25,6 +25,12 @@ class PydanticInsertArgumentsQuickFixTest : PydanticTestCase() {
     fun testPartArguments() {
         doTest(false)
     }
+    fun testNestedPartArguments() {
+        doTest(false)
+    }
+    fun testNestedOtherObject() {
+        doTest(false)
+    }
     fun testPartArgumentsOnlyRequired() {
         doTest(true)
     }