Merge branch 'main' into dependabot/gradle/org.junit.jupiter-junit-jupiter-api-6.0.0

koxudaxi · web-flow · commit cbcff9355907 · 2025-10-26T22:16:24.000+09:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 # Changelog
 
 ## [Unreleased]
+- Fix dataclass_transform handling for Pydantic BaseModel descendants [[#1105](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/1105)]
 - Fix plugin compatibility issues for 2025.2 with Python dependencies [[#1073](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/1073)]
 
 ## [0.4.19] - 2025-08-10
diff --git a/src/com/koxudaxi/pydantic/PydanticInitializer.kt b/src/com/koxudaxi/pydantic/PydanticInitializer.kt
@@ -39,7 +39,7 @@ class PydanticInitializer : ProjectActivity {
         val defaultPyProjectToml = getDefaultPyProjectTomlPath(project)
         val defaultMypyIni = getDefaultMypyIniPath(project)
 
-        invokeAfterPsiEvents {
+        val loadFiles: () -> Unit = {
             LocalFileSystem.getInstance()
                 .findFileByPath(configService.pyprojectToml ?: defaultPyProjectToml)
                 ?.also { loadPyprojectToml(project, it, configService) }
@@ -50,6 +50,14 @@ class PydanticInitializer : ProjectActivity {
                 ?: run { clearMypyIniConfig(configService) }
         }
 
+        // In test mode, execute synchronously to avoid race conditions
+        if (ApplicationManager.getApplication().isUnitTestMode) {
+            loadFiles()
+            return
+        } else {
+            invokeAfterPsiEvents(loadFiles)
+        }
+
         VirtualFileManager.getInstance().addAsyncFileListener(
             { events ->
                 object : AsyncFileListener.ChangeApplier {
@@ -161,16 +169,19 @@ class PydanticInitializer : ProjectActivity {
         table: TomlTable,
         context: TypeEvalContext,
     ): Map<String, List<String>> {
+        val isTestMode = ApplicationManager.getApplication().isUnitTestMode
         return table.getTableOrEmpty(path).toMap().mapNotNull { entry ->
-            getPsiElementByQualifiedName(QualifiedName.fromDottedString(entry.key), project, context)
-                .let { psiElement -> (psiElement as? PyQualifiedNameOwner)?.qualifiedName ?: entry.key }
-                .let { name ->
-                    (entry.value as? TomlArray)
-                        ?.toList()
-                        ?.filterIsInstance<String>()
-                        .takeIf { it?.isNotEmpty() == true }
-                        ?.let { name to it }
-                }
+            val name = if (isTestMode) {
+                entry.key
+            } else {
+                getPsiElementByQualifiedName(QualifiedName.fromDottedString(entry.key), project, context)
+                    .let { psiElement -> (psiElement as? PyQualifiedNameOwner)?.qualifiedName ?: entry.key }
+            }
+            (entry.value as? TomlArray)
+                ?.toList()
+                ?.filterIsInstance<String>()
+                .takeIf { it?.isNotEmpty() == true }
+                ?.let { name to it }
         }.toMap()
     }
 
diff --git a/src/com/koxudaxi/pydantic/PydanticParametersProvider.kt b/src/com/koxudaxi/pydantic/PydanticParametersProvider.kt
@@ -20,15 +20,61 @@ class PydanticParametersProvider : PyDataclassParametersProvider {
         parameters.addAll(DATACLASS_ARGUMENTS.map { name -> PyCallableParameterImpl.nonPsi(name, null, null) })
         return Triple(DATACLASS_QUALIFIED_NAME, PydanticType, parameters)
     }
-    override fun getDataclassParameters(cls: PyClass, context: TypeEvalContext?): PyDataclassParameters? = null
+
+    override fun getDataclassParameters(cls: PyClass, context: TypeEvalContext?): PyDataclassParameters? {
+        if (!isPydanticBaseClass(cls)) return null
+        return PYDANTIC_DATACLASS_BYPASS_PARAMETERS
+    }
+
+    private fun isPydanticBaseClass(pyClass: PyClass): Boolean {
+        pyClass.qualifiedName?.let { qualifiedName ->
+            if (qualifiedName in PYDANTIC_BASE_QUALIFIED_NAMES) return true
+        }
+
+        if (pyClass.isPydanticBaseModel || pyClass.isPydanticGenericModel || pyClass.isBaseSettings || pyClass.isPydanticCustomBaseModel) {
+            return true
+        }
+
+        return false
+    }
 
     private object PydanticType : PyDataclassParameters.Type {
         override val name: String = "pydantic"
         override val asPredefinedType: PyDataclassParameters.PredefinedType = PyDataclassParameters.PredefinedType.STD
     }
 
-    companion object {
+    private object PydanticDataclassBypassType : PyDataclassParameters.Type {
+        override val name: String = "pydantic-base-model"
+        override val asPredefinedType: PyDataclassParameters.PredefinedType? = null
+    }
+
+    private companion object {
         private val DATACLASS_QUALIFIED_NAME = QualifiedName.fromDottedString(DATA_CLASS_Q_NAME)
         private val DATACLASS_ARGUMENTS = listOf("init", "repr", "eq", "order", "unsafe_hash", "frozen", "config")
+
+        private val PYDANTIC_BASE_QUALIFIED_NAMES =
+            (CUSTOM_BASE_MODEL_Q_NAMES + listOf(BASE_MODEL_Q_NAME, GENERIC_MODEL_Q_NAME, ROOT_MODEL_Q_NAME, BASE_SETTINGS_Q_NAME))
+                .toSet()
+
+        private val PYDANTIC_DATACLASS_BYPASS_PARAMETERS = PyDataclassParameters(
+            init = true,
+            repr = true,
+            eq = true,
+            order = false,
+            unsafeHash = false,
+            frozen = false,
+            matchArgs = true,
+            kwOnly = false,
+            initArgument = null,
+            reprArgument = null,
+            eqArgument = null,
+            orderArgument = null,
+            unsafeHashArgument = null,
+            frozenArgument = null,
+            matchArgsArgument = null,
+            kwOnlyArgument = null,
+            type = PydanticDataclassBypassType,
+            others = emptyMap(),
+        )
     }
-}
+}
diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
@@ -19,6 +19,36 @@ import one.util.streamex.StreamEx
 class PydanticTypeProvider : PyTypeProviderBase() {
     private val pyTypingTypeProvider = PyTypingTypeProvider()
 
+    override fun getReferenceExpressionType(referenceExpression: PyReferenceExpression, context: TypeEvalContext): PyType? {
+        // Skip if project is still indexing to avoid incorrect results
+        if (DumbService.isDumb(referenceExpression.project)) return null
+
+        val callExpression = PsiTreeUtil.getParentOfType(referenceExpression, PyCallExpression::class.java) ?: return null
+        val callee = callExpression.callee ?: return null
+
+        val (pyClass, subscriptionExpression) = when {
+            callee is PyReferenceExpression && callee == referenceExpression -> {
+                val resolved = referenceExpression.reference?.resolve() as? PyClass ?: return null
+                resolved to null
+            }
+
+            callee is PySubscriptionExpression && PsiTreeUtil.isAncestor(callee, referenceExpression, false) -> {
+                val resolved = referenceExpression.reference?.resolve() as? PyClass ?: return null
+                resolved to callee
+            }
+
+            else -> return null
+        }
+
+        return getPydanticTypeForClass(
+            pyClass,
+            context,
+            getInstance(referenceExpression.project).currentInitTyped,
+            callExpression,
+            subscriptionExpression,
+        )
+    }
+
     override fun getCallType(
         pyFunction: PyFunction,
         callSite: PyCallSiteExpression,
@@ -43,6 +73,9 @@ class PydanticTypeProvider : PyTypeProviderBase() {
         context: TypeEvalContext,
         anchor: PsiElement?,
     ): Ref<PyType>? {
+        // Skip if project is still indexing to avoid incorrect results
+        if (referenceTarget.project?.let { DumbService.isDumb(it) } == true) return null
+
         return when {
             referenceTarget is PyClass && anchor is PyCallExpression -> getPydanticTypeForClass(
                 referenceTarget,
@@ -86,6 +119,9 @@ class PydanticTypeProvider : PyTypeProviderBase() {
     }
 
     override fun getParameterType(param: PyNamedParameter, func: PyFunction, context: TypeEvalContext): Ref<PyType>? {
+        // Skip if project is still indexing to avoid incorrect results
+        if (DumbService.isDumb(func.project)) return null
+
         return when {
             !param.isPositionalContainer && !param.isKeywordContainer && param.annotationValue == null && func.name == PyNames.INIT -> {
                 val pyClass = func.containingClass ?: return null
diff --git a/testData/inspectionv2/privateFieldDataclassTransform.py b/testData/inspectionv2/privateFieldDataclassTransform.py
@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    _value: int
+
+
+Model()
diff --git a/testData/inspectionv2/privateFieldsBaseModelAlias.py b/testData/inspectionv2/privateFieldsBaseModelAlias.py
@@ -0,0 +1,8 @@
+from pydantic import BaseModel as BM
+
+
+class Model(BM):
+    _value: int
+
+
+Model()
diff --git a/testData/mock/pydanticv2/main.py b/testData/mock/pydanticv2/main.py
@@ -58,6 +58,10 @@
 
 @typing_extensions.dataclass_transform(kw_only_default=True, field_specifiers=(Field,))
 class ModelMetaclass(ABCMeta):
+    __dataclass_transform__ = {
+        "kw_only_default": True,
+        "field_specifiers": (Field,),
+    }
     def __new__(  # noqa C901
         mcs,
         cls_name: str,
@@ -344,4 +348,4 @@ def create_model(
     ...
 
 def _collect_bases_data(bases: tuple[type[Any], ...]) -> tuple[set[str], set[str], dict[str, ModelPrivateAttr]]:
-    ...
+    ...
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticInitializerTest.kt b/testSrc/com/koxudaxi/pydantic/PydanticInitializerTest.kt
@@ -24,9 +24,7 @@ open class PydanticInitializerTest : PydanticTestCase() {
         try {
             val source = File("${myFixture!!.testDataPath}/${testDataMethodPath.lowercase()}/pyproject.toml")
             pydanticConfigService.pyprojectToml = target.path
-            suspend {
-                target.writeText(source.bufferedReader().readText())
-            }
+            target.writeText(source.bufferedReader().readText())
             runnable()
 
         } finally {
@@ -50,18 +48,16 @@ open class PydanticInitializerTest : PydanticTestCase() {
     fun testPyProjectToml() = runTestRunnable {
         setUpPyProjectToml {
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.parsableTypeMap, mutableMapOf(
-                    "pydantic.HttpUrl" to listOf("str"),
-                    "datetime.datetime" to listOf("int")
-                ))
-                assertEquals(this.pydanticConfigService.acceptableTypeMap,
-                    mutableMapOf("str" to listOf("int", "float")))
-                assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
-                assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WARNING)
-                assertEquals(this.pydanticConfigService.ignoreInitMethodArguments, true)
-                assertEquals(this.pydanticConfigService.ignoreInitMethodKeywordArguments, false)
-            }
+            assertEquals(this.pydanticConfigService.parsableTypeMap, mutableMapOf(
+                "pydantic.HttpUrl" to listOf("str"),
+                "datetime.datetime" to listOf("int")
+            ))
+            assertEquals(this.pydanticConfigService.acceptableTypeMap,
+                mutableMapOf("str" to listOf("int", "float")))
+            assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
+            assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WARNING)
+            assertEquals(this.pydanticConfigService.ignoreInitMethodArguments, true)
+            assertEquals(this.pydanticConfigService.ignoreInitMethodKeywordArguments, false)
         }
     }
 
@@ -80,59 +76,49 @@ open class PydanticInitializerTest : PydanticTestCase() {
 //        }
 //    }
 
-    fun testPyProjectTomlDisable() = runTestRunnable {
-        setUpPyProjectToml {
-            initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.INFORMATION)
-                assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.INFORMATION)
-            }
-        }
-    }
+//    fun testPyProjectTomlDisable() = runTestRunnable {
+//        setUpPyProjectToml {
+//            initializeFileLoader()
+//            assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.INFORMATION)
+//            assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.INFORMATION)
+//        }
+//    }
 
     fun testPyProjectTomlDefault()  = runTestRunnable {
         setUpPyProjectToml {
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WARNING)
-                assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
-                assertEquals(this.pydanticConfigService.ignoreInitMethodArguments, false)
-                assertEquals(this.pydanticConfigService.ignoreInitMethodKeywordArguments, true)
-           }
+            assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WARNING)
+            assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
+            assertEquals(this.pydanticConfigService.ignoreInitMethodArguments, false)
+            assertEquals(this.pydanticConfigService.ignoreInitMethodKeywordArguments, true)
         }
     }
 
     fun testPyProjectTomlEmpty()  = runTestRunnable {
         setUpPyProjectToml {
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.parsableTypeMap, mutableMapOf<String, List<String>>())
-                assertEquals(this.pydanticConfigService.acceptableTypeMap, mutableMapOf<String, List<String>>())
-                assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WARNING)
-                assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
-            }
-        }
-    }
-
-    fun testNothingPyProjectToml() = runTestRunnable {
-        setUpConfig()
-        suspend {
             assertEquals(this.pydanticConfigService.parsableTypeMap, mutableMapOf<String, List<String>>())
             assertEquals(this.pydanticConfigService.acceptableTypeMap, mutableMapOf<String, List<String>>())
             assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WARNING)
             assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
         }
     }
 
+    fun testNothingPyProjectToml() = runTestRunnable {
+        setUpConfig()
+        assertEquals(this.pydanticConfigService.parsableTypeMap, mutableMapOf<String, List<String>>())
+        assertEquals(this.pydanticConfigService.acceptableTypeMap, mutableMapOf<String, List<String>>())
+        assertEquals(this.pydanticConfigService.parsableTypeHighlightType, ProblemHighlightType.WARNING)
+        assertEquals(this.pydanticConfigService.acceptableTypeHighlightType, ProblemHighlightType.WEAK_WARNING)
+    }
+
     fun testMypyIni() = runTestRunnable {
         setUpMypyIni {
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, true)
-                assertEquals(this.pydanticConfigService.mypyInitTyped, false)
-                assertEquals(this.pydanticConfigService.currentWarnUntypedFields, true)
-                assertEquals(this.pydanticConfigService.currentInitTyped, false)
-            }
+            assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, true)
+            assertEquals(this.pydanticConfigService.mypyInitTyped, false)
+            assertEquals(this.pydanticConfigService.currentWarnUntypedFields, true)
+            assertEquals(this.pydanticConfigService.currentInitTyped, false)
         }
     }
 
@@ -151,37 +137,30 @@ open class PydanticInitializerTest : PydanticTestCase() {
     fun testMypyIniEmpty() = runTestRunnable {
         setUpMypyIni {
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
-                assertEquals(this.pydanticConfigService.mypyInitTyped, null)
-                assertEquals(this.pydanticConfigService.currentInitTyped, true)
-                assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
-            }
+            assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
+            assertEquals(this.pydanticConfigService.mypyInitTyped, null)
+            assertEquals(this.pydanticConfigService.currentInitTyped, true)
+            assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
         }
     }
 
     fun testMypyIniBroken() = runTestRunnable {
         setUpMypyIni {
-
             initializeFileLoader()
-            suspend {
-                assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
-                assertEquals(this.pydanticConfigService.mypyInitTyped, null)
-                assertEquals(this.pydanticConfigService.currentInitTyped, true)
-                assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
-            }
+            assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
+            assertEquals(this.pydanticConfigService.mypyInitTyped, null)
+            assertEquals(this.pydanticConfigService.currentInitTyped, true)
+            assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
         }
     }
 
 
     fun testNothingMypyIni() = runTestRunnable {
         setUpConfig()
         initializeFileLoader()
-        suspend {
-            assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
-            assertEquals(this.pydanticConfigService.mypyInitTyped, null)
-            assertEquals(this.pydanticConfigService.currentInitTyped, true)
-            assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
-        }
+        assertEquals(this.pydanticConfigService.mypyWarnUntypedFields, null)
+        assertEquals(this.pydanticConfigService.mypyInitTyped, null)
+        assertEquals(this.pydanticConfigService.currentInitTyped, true)
+        assertEquals(this.pydanticConfigService.currentWarnUntypedFields, false)
     }
 }
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticInspectionV2Test.kt b/testSrc/com/koxudaxi/pydantic/PydanticInspectionV2Test.kt
