Merge pull request #34 from koxudaxi/support_refactoring_fields_by_ke…

…yword_argument support refactoring fields by a keyword argument
koxudaxi · Aug 14, 2019 · 79f62d4 · 79f62d4
2 parents 6caf5ab + 9748a59
commit 79f62d4
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 38 deletions.
diff --git a/README.md b/README.md
@@ -12,7 +12,7 @@ A JetBrains PyCharm plugin for [`pydantic`](https://github.com/samuelcolvin/pyda
 * Model-specific `__init__`-signature inspection and autocompletion for subclasses of `pydantic.BaseModel`
 * 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, PyCharm will present a dialog offering the choice to automatically rename the keyword where it occurs in a model initialization call.
+  * (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.
 
 
 ## How to install:

diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
@@ -1,9 +1,18 @@
 <idea-plugin url="https://github.com/koxudaxi/pydantic-pycharm-plugin">
     <id>com.koxudaxi.pydantic</id>
     <name>Pydantic</name>
-    <version>0.0.11</version>
+    <version>0.0.12</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
-
+    <change-notes><![CDATA[
+    <h2>version 0.0.12</h2>
+    <p>Features</p>
+    <ul>
+        <li>Support refactoring fields by a keyword argument [#34] </li>
+        <li>Support refactoring super-classes and inheritor-classes [#34] </li>
+        <li>Support ellipsis(...) in fields [#34] </li>
+        <li>Support Schema in fields [#31] </li>
+    </ul>
+    ]]></change-notes>
     <description><![CDATA[
     <p>This plugin provides autocompletion support for <a href="https://github.com/samuelcolvin/pydantic">pydantic</a></p>
     <h2>Features</h2>
@@ -13,7 +22,7 @@
                <li>Model-specific __init__-signature inspection and autocompletion for subclasses of pydantic.BaseModel</li>
                <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, PyCharm will present a dialog offering the choice to automatically rename the keyword where it occurs in a model initialization call.</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>
             </ul>
          </li>
         <li>pydantic.dataclasses.dataclass

diff --git a/src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt b/src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt
@@ -1,24 +1,32 @@
 package com.koxudaxi.pydantic
 
 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.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 java.util.*
+import com.jetbrains.python.psi.search.PyClassInheritorsSearch
+import com.jetbrains.python.psi.types.TypeEvalContext
+
 
 class PydanticFieldRenameFactory : AutomaticRenamerFactory {
     override fun isApplicable(element: PsiElement): Boolean {
-        // Field to KeywordArguments
-        if (element is PyTargetExpression) {
-            val pyClass = element.containingClass ?: return false
-            if (pyClass.isSubclass("pydantic.main.BaseModel", null)) return true
+        when (element) {
+            is PyTargetExpression -> {
+                val pyClass = element.containingClass ?: return false
+                if (pyClass.isSubclass("pydantic.main.BaseModel", null)) return true
+            }
+            is PyKeywordArgument -> {
+                val pyClass = getPyClassByPyKeywordArgument(element) ?: return false
+                if (pyClass.isSubclass("pydantic.main.BaseModel", null)) return true
+            }
         }
         return false
     }
@@ -40,41 +48,81 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
     }
 
     class PydanticFieldRenamer(element: PsiElement, newName: String) : AutomaticRenamer() {
-
         init {
-            val pyTargetExpression = (element as? PyTargetExpression)
-            if (pyTargetExpression?.name != null) {
-
-                val pyClass = pyTargetExpression.containingClass
-                ReferencesSearch.search(pyClass as @org.jetbrains.annotations.NotNull PsiElement).forEach { psiReference ->
-                    val callee = PsiTreeUtil.getParentOfType(psiReference.element, PyCallExpression::class.java)
-                    callee?.arguments?.forEach { argument ->
-                        if (argument is PyKeywordArgument) {
-                            if (argument.name == pyTargetExpression.name) {
-                                myElements.add(argument)
-                            }
-                        }
+            val added = mutableSetOf<PyClass>()
+            when (element) {
+                is PyTargetExpression -> if (element.name != null) {
+                    val pyClass = element.containingClass
+                    addAllElement(pyClass, element.name!!, added)
+                    suggestAllNames(element.name!!, newName)
+                }
+                is PyKeywordArgument -> if (element.name != null) {
+                    val pyClass = getPyClassByPyKeywordArgument(element)
+                    addAllElement(pyClass, element.name!!, added)
+                    suggestAllNames(element.name!!, newName)
+                }
+            }
+        }
+
+        private fun addAllElement(pyClass: PyClass?, elementName: String, added: MutableSet<PyClass>) {
+            if (pyClass == null) return
+            added.add(pyClass)
+            addClassAttributes(pyClass, elementName)
+            addKeywordArguments(pyClass, elementName)
+            pyClass.getAncestorClasses(null).forEach {  ancestorClass ->
+                if (ancestorClass.qualifiedName != "pydantic.main.BaseModel") {
+                    if (ancestorClass.isSubclass("pydantic.main.BaseModel", null) &&
+                            !added.contains(ancestorClass)) {
+                        addAllElement(ancestorClass, elementName, added)
                     }
-                    return@forEach
                 }
-            suggestAllNames(element.name!!, newName)
+            }
+            PyClassInheritorsSearch.search(pyClass, true).forEach { inheritorsPyClass ->
+                if (inheritorsPyClass.qualifiedName != "pydantic.main.BaseModel" && ! added.contains(inheritorsPyClass)) {
+                    addAllElement(inheritorsPyClass, elementName, added)
+                }
+            }
         }
-    }
 
-    override fun getDialogTitle(): String {
-        return "Rename Fields"
-    }
+        private fun addClassAttributes(pyClass: PyClass, elementName: String) {
+            pyClass.classAttributes.forEach { pyTargetExpression ->
+                if (pyTargetExpression.name == elementName) {
+                    myElements.add(pyTargetExpression)
+                }
+            }
+        }
 
-    override fun getDialogDescription(): String {
-        return "Rename field in hierarchy to:"
-    }
+        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)
 
-    override fun entityName(): String {
-        return "Field"
-    }
+                    }
+                }
+            }
+        }
+        override fun getDialogTitle(): String {
+            return "Rename Fields"
+        }
+
+        override fun getDialogDescription(): String {
+            return "Rename field in hierarchy to:"
+        }
 
-    override fun isSelectedByDefault(): Boolean {
-        return true
+        override fun entityName(): String {
+            return "Field"
+        }
+
+        override fun isSelectedByDefault(): Boolean {
+            return true
+        }
     }
+
 }
-}
+
+private fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument) : PyClass? {
+    val pyCallExpression = pyKeywordArgument.parent?.parent as? PyCallExpression ?: return null
+    return pyCallExpression.callee?.reference?.resolve() as? PyClass ?: return null
+}