Support to search related fields

README.md

@@ -18,6 +18,7 @@
 * Model-specific `__init__`-arguments type-checking for subclasses of `pydantic.BaseModel` 
 * Refactor support for renaming fields for subclasses of `BaseModel`
   * (If the field name is refactored from the model definition or `__init__` call keyword arguments, PyCharm will present a dialog offering the choice to automatically rename the keyword where it occurs in a model initialization call.
+* Search related-fields by class attributes and keyword arguments of `__init__` with `Ctrl+B` and `Cmd+B`
 
 
 ## How to install:

build.gradle

@@ -36,16 +36,16 @@ allprojects {
     compileKotlin {
         kotlinOptions {
             jvmTarget = "1.8"
-            languageVersion = "1.2"
-            apiVersion = "1.2"
+            languageVersion = "1.3"
+            apiVersion = "1.3"
         }
     }
 
     compileTestKotlin {
         kotlinOptions {
             jvmTarget = "1.8"
-            languageVersion = "1.2"
-            apiVersion = "1.2"
+            languageVersion = "1.3"
+            apiVersion = "1.23"
         }
     }
 

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.0.13</version>
+    <version>0.0.14</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.0.14</h2>
+    <p>Features</p>
+    <ul>
+        <li>Search related-fields by class attributes and keyword arguments of __init__. with Ctrl+B and Cmd+B [#42] </li>
+    </ul>
     <h2>version 0.0.13</h2>
     <p>Features, BugFixes</p>
     <ul>
@@ -29,6 +34,7 @@
                <li>Model-specific __init__-arguments type-checking for subclasses of pydantic.BaseModel</li>
                <li>Refactor support for renaming fields for subclasses of BaseModel</li>
                 <li>(If the field name is refactored from the model definition or __init__ call keyword arguments, PyCharm will present a dialog offering the choice to automatically rename the keyword where it occurs in a model initialization call.</li>
+               <li>Search related-fields by class attributes and keyword arguments of __init__. with Ctrl+B and Cmd+B</li>
             </ul>
          </li>
         <li>pydantic.dataclasses.dataclass
@@ -50,6 +56,7 @@
                          enabledByDefault="true" level="WARNING"
                          implementationClass="com.koxudaxi.pydantic.PydanticInspection"/>
         <automaticRenamerFactory implementation="com.koxudaxi.pydantic.PydanticFieldRenameFactory"/>
+        <referencesSearch implementation="com.koxudaxi.pydantic.PydanticFieldSearchExecutor"/>
     </extensions>
     <extensions defaultExtensionNs="Pythonid">
         <typeProvider implementation="com.koxudaxi.pydantic.PydanticTypeProvider"/>

src/com/koxudaxi/pydantic/PydanticBaseModel.kt

@@ -0,0 +1,11 @@
+package com.koxudaxi.pydantic
+
+import com.jetbrains.python.psi.PyCallExpression
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyKeywordArgument
+
+
+fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument) : PyClass? {
+    val pyCallExpression = pyKeywordArgument.parent?.parent as? PyCallExpression ?: return null
+    return pyCallExpression.callee?.reference?.resolve() as? PyClass ?: return null
+}

src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt

@@ -6,14 +6,12 @@ import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.refactoring.rename.naming.AutomaticRenamer
 import com.intellij.refactoring.rename.naming.AutomaticRenamerFactory
 import com.intellij.usageView.UsageInfo
-import com.jetbrains.extensions.python.inherits
 import com.jetbrains.python.codeInsight.PyCodeInsightSettings
 import com.jetbrains.python.psi.PyCallExpression
 import com.jetbrains.python.psi.PyClass
 import com.jetbrains.python.psi.PyKeywordArgument
 import com.jetbrains.python.psi.PyTargetExpression
 import com.jetbrains.python.psi.search.PyClassInheritorsSearch
-import com.jetbrains.python.psi.types.TypeEvalContext
 
 
 class PydanticFieldRenameFactory : AutomaticRenamerFactory {
@@ -69,7 +67,7 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
             added.add(pyClass)
             addClassAttributes(pyClass, elementName)
             addKeywordArguments(pyClass, elementName)
-            pyClass.getAncestorClasses(null).forEach {  ancestorClass ->
+            pyClass.getAncestorClasses(null).forEach { ancestorClass ->
                 if (ancestorClass.qualifiedName != "pydantic.main.BaseModel") {
                     if (ancestorClass.isSubclass("pydantic.main.BaseModel", null) &&
                             !added.contains(ancestorClass)) {
@@ -78,31 +76,29 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
                 }
             }
             PyClassInheritorsSearch.search(pyClass, true).forEach { inheritorsPyClass ->
-                if (inheritorsPyClass.qualifiedName != "pydantic.main.BaseModel" && ! added.contains(inheritorsPyClass)) {
+                if (inheritorsPyClass.qualifiedName != "pydantic.main.BaseModel" && !added.contains(inheritorsPyClass)) {
                     addAllElement(inheritorsPyClass, elementName, added)
                 }
             }
         }
 
         private fun addClassAttributes(pyClass: PyClass, elementName: String) {
-            pyClass.classAttributes.forEach { pyTargetExpression ->
-                if (pyTargetExpression.name == elementName) {
-                    myElements.add(pyTargetExpression)
-                }
-            }
+            val pyTargetExpression = pyClass.findClassAttribute(elementName, false, null) ?: return
+            myElements.add(pyTargetExpression)
         }
 
         private fun addKeywordArguments(pyClass: PyClass, elementName: String) {
             ReferencesSearch.search(pyClass as PsiElement).forEach { psiReference ->
                 val callee = PsiTreeUtil.getParentOfType(psiReference.element, PyCallExpression::class.java)
                 callee?.arguments?.forEach { argument ->
                     if (argument is PyKeywordArgument && argument.name == elementName) {
-                            myElements.add(argument)
+                        myElements.add(argument)
 
                     }
                 }
             }
         }
+
         override fun getDialogTitle(): String {
             return "Rename Fields"
         }
@@ -121,8 +117,3 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
     }
 
 }
-
-private fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument) : PyClass? {
-    val pyCallExpression = pyKeywordArgument.parent?.parent as? PyCallExpression ?: return null
-    return pyCallExpression.callee?.reference?.resolve() as? PyClass ?: return null
-}

src/com/koxudaxi/pydantic/PydanticFieldSearchExecutor.kt

@@ -0,0 +1,86 @@
+package com.koxudaxi.pydantic
+
+import com.intellij.openapi.application.QueryExecutorBase
+import com.intellij.openapi.application.ReadAction.run
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiReference
+import com.intellij.psi.search.searches.ReferencesSearch
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.Processor
+import com.jetbrains.python.psi.PyCallExpression
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyKeywordArgument
+import com.jetbrains.python.psi.PyTargetExpression
+import com.jetbrains.python.psi.search.PyClassInheritorsSearch
+
+private fun searchField(pyClass: PyClass, elementName: String, consumer: Processor<in PsiReference>): Boolean {
+    if (!pyClass.isSubclass("pydantic.main.BaseModel", null)) return false
+    val pyTargetExpression = pyClass.findClassAttribute(elementName, false, null) ?: return false
+    consumer.process(pyTargetExpression.reference)
+    return true
+}
+
+private fun searchKeywordArgument(pyClass: PyClass, elementName: String, consumer: Processor<in PsiReference>) {
+    if (!pyClass.isSubclass("pydantic.main.BaseModel", null)) return
+    ReferencesSearch.search(pyClass as PsiElement).forEach { psiReference ->
+        val callee = PsiTreeUtil.getParentOfType(psiReference.element, PyCallExpression::class.java)
+        callee?.arguments?.forEach { argument ->
+            if (argument is PyKeywordArgument && argument.name == elementName) {
+                consumer.process(argument.reference)
+
+            }
+        }
+    }
+}
+
+private fun searchDirectReferenceField(pyClass: PyClass, elementName: String, consumer: Processor<in PsiReference>): Boolean {
+    if (searchField(pyClass, elementName, consumer)) return true
+
+    pyClass.getAncestorClasses(null).forEach { ancestorClass ->
+        if (ancestorClass.qualifiedName != "pydantic.main.BaseModel") {
+            if (ancestorClass.isSubclass("pydantic.main.BaseModel", null)) {
+                if (searchDirectReferenceField(ancestorClass, elementName, consumer)) {
+                    return true
+                }
+            }
+        }
+    }
+    return false
+}
+
+private fun searchAllElementReference(pyClass: PyClass?, elementName: String, added: MutableSet<PyClass>, consumer: Processor<in PsiReference>) {
+    if (pyClass == null) return
+    added.add(pyClass)
+    searchField(pyClass, elementName, consumer)
+    searchKeywordArgument(pyClass, elementName, consumer)
+    pyClass.getAncestorClasses(null).forEach {  ancestorClass ->
+        if (ancestorClass.qualifiedName != "pydantic.main.BaseModel" &&  !added.contains(ancestorClass)){
+            searchField(pyClass, elementName, consumer)
+        }
+    }
+    PyClassInheritorsSearch.search(pyClass, true).forEach { inheritorsPyClass ->
+        if (inheritorsPyClass.qualifiedName != "pydantic.main.BaseModel" && !added.contains(inheritorsPyClass)) {
+            searchAllElementReference(inheritorsPyClass, elementName, added, consumer)
+        }
+    }
+}
+
+class PydanticFieldSearchExecutor : QueryExecutorBase<PsiReference, ReferencesSearch.SearchParameters>() {
+    override fun processQuery(queryParameters: ReferencesSearch.SearchParameters, consumer: Processor<in PsiReference>) {
+
+        when (val element = queryParameters.elementToSearch) {
+            is PyKeywordArgument -> run<RuntimeException> {
+                val elementName = element.name ?: return@run
+                val pyClass = getPyClassByPyKeywordArgument(element) ?: return@run
+                if (!pyClass.isSubclass("pydantic.main.BaseModel", null)) return@run
+                searchDirectReferenceField(pyClass, elementName, consumer)
+            }
+            is PyTargetExpression -> run<RuntimeException> {
+                val elementName = element.name ?: return@run
+                val pyClass = element.containingClass ?: return@run
+                if (!pyClass.isSubclass("pydantic.main.BaseModel", null)) return@run
+                searchAllElementReference(pyClass, elementName, mutableSetOf(), consumer)
+            }
+        }
+    }
+}