From c3d6215413969c0f4ff65a4c8341148eae2629c9 Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Thu, 27 Jun 2024 19:19:43 +0000
Subject: [PATCH 1/5] Add gritql runtime_data migration rule

---
 .grit/.gitignore                         |  2 +
 .grit/grit.yaml                          |  3 +
 .grit/patterns/runtime_data_migration.md | 71 ++++++++++++++++++++++++
 3 files changed, 76 insertions(+)
 create mode 100644 .grit/.gitignore
 create mode 100644 .grit/grit.yaml
 create mode 100644 .grit/patterns/runtime_data_migration.md

diff --git a/.grit/.gitignore b/.grit/.gitignore
new file mode 100644
index 00000000000..e4fdfb17c16
--- /dev/null
+++ b/.grit/.gitignore
@@ -0,0 +1,2 @@
+.gritmodules*
+*.log
diff --git a/.grit/grit.yaml b/.grit/grit.yaml
new file mode 100644
index 00000000000..865cd4367e9
--- /dev/null
+++ b/.grit/grit.yaml
@@ -0,0 +1,3 @@
+version: 0.0.1
+patterns:
+#  - name: github.com/getgrit/stdlib#*
diff --git a/.grit/patterns/runtime_data_migration.md b/.grit/patterns/runtime_data_migration.md
new file mode 100644
index 00000000000..5bdb48c5334
--- /dev/null
+++ b/.grit/patterns/runtime_data_migration.md
@@ -0,0 +1,71 @@
+---
+tags: [migration, code_quality]
+---
+# Migrate integration from hass.data to entry.runtime_data
+
+Migrate an integration from hass.data to entry.runtime_data
+
+```grit
+engine marzano(0.1)
+language python
+
+pattern refactor_functions($config_entry_type) {
+    function_definition($parameters, $body) where {
+        // change entry type
+        $entry = $parameters[1],
+        $entry <: typed_parameter(name=$entry_name, $type) where {
+            $type <: type(type="ConfigEntry"),
+            $type => $config_entry_type
+        },
+
+        // migrate hass.data to entry.runtime_data
+        $body <: maybe contains assignment($left, right=$runtime_data) as $assignment where {
+            $runtime_data <: `hass.data[$_][entry.entry_id]`,
+            $assignment => `$left = $entry_name.runtime_data`
+        },
+    }
+}
+
+pattern refactor_init($config_entry_type) {
+    function_definition(name="async_setup_entry", $parameters, $body) as $func where {
+        // change entry type
+        $entry = $parameters[1],
+        $entry <: typed_parameter(name=$entry_name, $type) where {
+            $type => $config_entry_type
+        },
+
+        // migrate hass.data to entry.runtime_data
+        $body <: contains or {
+            `hass.data.setdefault($...)[entry.entry_id]`,
+            `hass.data[$_][entry.entry_id]`,
+            } as $runtime_data where {
+                $runtime_data => `$entry_name.runtime_data`
+            },
+
+        $config_entry_type_definition = `# TODO: Please add the correct type\n`,
+        $config_entry_type_definition += `type $config_entry_type = ConfigEntry`,
+        $func => `$config_entry_type_definition\n\n$func`
+    }
+}
+
+multifile {
+    bubble($domain, $config_entry_type) file($name, $body) where {
+        $file_parts = split($name, "/"),
+        $components_folder = $file_parts[-3],
+        $components_folder <: includes `components`, // with includes we allow also custom_components
+        $domain = $file_parts[-2],
+        $config_entry_type = capitalize($domain),
+        $config_entry_type += "ConfigEntry",
+        $name <: includes `__init__.py`,
+        $body <: contains and {
+            refactor_init($config_entry_type),
+            maybe refactor_functions($config_entry_type)
+        },
+    },
+    bubble($domain, $config_entry_type) file($name, $body) where {
+        $file_parts = split($name, "/"),
+        $domain = $file_parts[-2],
+        $body <: contains refactor_functions($config_entry_type)
+    }
+}
+```
\ No newline at end of file

From ae1339d1ed54aa2898f87cbaf1f51067dfcae02d Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Mon, 1 Jul 2024 07:53:23 +0000
Subject: [PATCH 2/5] Another try

---
 .grit/patterns/runtime_data_migration.md | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/.grit/patterns/runtime_data_migration.md b/.grit/patterns/runtime_data_migration.md
index 5bdb48c5334..6d53595534b 100644
--- a/.grit/patterns/runtime_data_migration.md
+++ b/.grit/patterns/runtime_data_migration.md
@@ -15,7 +15,8 @@ pattern refactor_functions($config_entry_type) {
         $entry = $parameters[1],
         $entry <: typed_parameter(name=$entry_name, $type) where {
             $type <: type(type="ConfigEntry"),
-            $type => $config_entry_type
+            $type => $config_entry_type,
+            //$config_entry_type <: ensure_import_from(source = `.`),
         },
 
         // migrate hass.data to entry.runtime_data
@@ -49,22 +50,32 @@ pattern refactor_init($config_entry_type) {
 }
 
 multifile {
-    bubble($domain, $config_entry_type) file($name, $body) where {
+    bubble($domain_list) file($name, $body) where {
+        $name <: includes `__init__.py`,
         $file_parts = split($name, "/"),
         $components_folder = $file_parts[-3],
         $components_folder <: includes `components`, // with includes we allow also custom_components
         $domain = $file_parts[-2],
         $config_entry_type = capitalize($domain),
         $config_entry_type += "ConfigEntry",
-        $name <: includes `__init__.py`,
         $body <: contains and {
             refactor_init($config_entry_type),
-            maybe refactor_functions($config_entry_type)
+            maybe refactor_functions($config_entry_type),
+            if ($domain_list <: undefined) {
+                $domain_list = []
+            },
+            $domain_list += $domain
         },
     },
-    bubble($domain, $config_entry_type) file($name, $body) where {
+    bubble($domain_list) file($name, $body) where {
+        $name <: not includes `__init__.py`,
         $file_parts = split($name, "/"),
+        $components_folder = $file_parts[-3],
+        $components_folder <: includes `components`, // with includes we allow also custom_components
         $domain = $file_parts[-2],
+        $domain_list <: includes $domain,
+        $config_entry_type = capitalize($domain),
+        $config_entry_type += "ConfigEntry",
         $body <: contains refactor_functions($config_entry_type)
     }
 }

From 855f79cc7219829462878c5e964b3d127ac89a7d Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Mon, 1 Jul 2024 12:57:23 +0000
Subject: [PATCH 3/5] improve

---
 .grit/patterns/runtime_data_migration.md | 103 +++++++++++------------
 1 file changed, 48 insertions(+), 55 deletions(-)

diff --git a/.grit/patterns/runtime_data_migration.md b/.grit/patterns/runtime_data_migration.md
index 6d53595534b..641ed43596b 100644
--- a/.grit/patterns/runtime_data_migration.md
+++ b/.grit/patterns/runtime_data_migration.md
@@ -6,77 +6,70 @@ tags: [migration, code_quality]
 Migrate an integration from hass.data to entry.runtime_data
 
 ```grit
-engine marzano(0.1)
 language python
 
-pattern refactor_functions($config_entry_type) {
-    function_definition($parameters, $body) where {
-        // change entry type
-        $entry = $parameters[1],
-        $entry <: typed_parameter(name=$entry_name, $type) where {
+pattern refactor_functions($config_entry_type, $file_name, $config_entry_type_defined) {
+    function_definition($parameters, $body) as $func where {
+        // change config entry type
+        $parameters <: contains typed_parameter(name=$entry_name, $type) where {
             $type <: type(type="ConfigEntry"),
             $type => $config_entry_type,
-            //$config_entry_type <: ensure_import_from(source = `.`),
-        },
-
-        // migrate hass.data to entry.runtime_data
-        $body <: maybe contains assignment($left, right=$runtime_data) as $assignment where {
-            $runtime_data <: `hass.data[$_][entry.entry_id]`,
-            $assignment => `$left = $entry_name.runtime_data`
-        },
-    }
-}
-
-pattern refactor_init($config_entry_type) {
-    function_definition(name="async_setup_entry", $parameters, $body) as $func where {
-        // change entry type
-        $entry = $parameters[1],
-        $entry <: typed_parameter(name=$entry_name, $type) where {
-            $type => $config_entry_type
-        },
-
-        // migrate hass.data to entry.runtime_data
-        $body <: contains or {
-            `hass.data.setdefault($...)[entry.entry_id]`,
-            `hass.data[$_][entry.entry_id]`,
-            } as $runtime_data where {
-                $runtime_data => `$entry_name.runtime_data`
+            if ($file_name <: not `__init__`) {
+                //$config_entry_type <: ensure_import_from(source = `.`),
             },
+        },
 
-        $config_entry_type_definition = `# TODO: Please add the correct type\n`,
-        $config_entry_type_definition += `type $config_entry_type = ConfigEntry`,
-        $func => `$config_entry_type_definition\n\n$func`
+        if (and {$file_name <: `__init__`, $config_entry_type_defined <: undefined}) {
+            if ($func <: within decorated_definition() as $decorated) {
+                // we need to insert the new type before all function decorators
+                $func = $decorated,
+            },
+            $config_entry_type_definition = `# TODO: Please add the correct type\n`,
+            $config_entry_type_definition += `type $config_entry_type = ConfigEntry`,
+            $func => `$config_entry_type_definition\n\n$func`,
+            $config_entry_type_defined = true,
+        },
+
+        // migrate hass.data to entry.runtime_data
+        $body <: maybe contains assignment($left, $right) as $assignment where {
+            or {
+                and {
+                    $right <: `hass.data[$_][entry.entry_id]`,
+                    $assignment => `$left = $entry_name.runtime_data`
+                },
+                and {
+                    $left <: or {
+                            `hass.data.setdefault($...)[entry.entry_id]`,
+                            `hass.data[$_][entry.entry_id]`,
+                        } as $runtime_data where {
+                            $runtime_data => `$entry_name.runtime_data`
+                        },
+                }
+            }
+        },
     }
 }
 
 multifile {
     bubble($domain_list) file($name, $body) where {
-        $name <: includes `__init__.py`,
-        $file_parts = split($name, "/"),
-        $components_folder = $file_parts[-3],
-        $components_folder <: includes `components`, // with includes we allow also custom_components
-        $domain = $file_parts[-2],
-        $config_entry_type = capitalize($domain),
-        $config_entry_type += "ConfigEntry",
-        $body <: contains and {
-            refactor_init($config_entry_type),
-            maybe refactor_functions($config_entry_type),
-            if ($domain_list <: undefined) {
-                $domain_list = []
-            },
-            $domain_list += $domain
+        // find all integrations, which can be migrated
+        $filename <: r".*components/([^/]+)/__init__\.py$"($domain),
+        $body <: contains or {
+        `hass.data.setdefault($...)[entry.entry_id]`,
+        `hass.data[$_][entry.entry_id]`,
         },
+        if ($domain_list <: undefined) {
+            $domain_list = []
+        },
+        $domain_list += $domain,
     },
     bubble($domain_list) file($name, $body) where {
-        $name <: not includes `__init__.py`,
-        $file_parts = split($name, "/"),
-        $components_folder = $file_parts[-3],
-        $components_folder <: includes `components`, // with includes we allow also custom_components
-        $domain = $file_parts[-2],
+        // migrate files
+        $filename <: r".*components/([^/]+)/([^/]+)\.py$"($domain, $file_name),
         $domain_list <: includes $domain,
         $config_entry_type = capitalize($domain),
         $config_entry_type += "ConfigEntry",
-        $body <: contains refactor_functions($config_entry_type)
-    }
+        $body <: contains refactor_functions($config_entry_type, $file_name, $config_entry_type_defined),
+    },
 }
 ```
\ No newline at end of file

From be316cafab5903e4ef6168b265f3d391ac7fdcb4 Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Tue, 2 Jul 2024 08:35:12 +0000
Subject: [PATCH 4/5] Add import

---
 .grit/grit.yaml                          | 2 +-
 .grit/patterns/runtime_data_migration.md | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/.grit/grit.yaml b/.grit/grit.yaml
index 865cd4367e9..c93f16d6ecd 100644
--- a/.grit/grit.yaml
+++ b/.grit/grit.yaml
@@ -1,3 +1,3 @@
 version: 0.0.1
 patterns:
-#  - name: github.com/getgrit/stdlib#*
+  - name: github.com/getgrit/stdlib#*
diff --git a/.grit/patterns/runtime_data_migration.md b/.grit/patterns/runtime_data_migration.md
index 641ed43596b..0c662a66fea 100644
--- a/.grit/patterns/runtime_data_migration.md
+++ b/.grit/patterns/runtime_data_migration.md
@@ -15,7 +15,7 @@ pattern refactor_functions($config_entry_type, $file_name, $config_entry_type_de
             $type <: type(type="ConfigEntry"),
             $type => $config_entry_type,
             if ($file_name <: not `__init__`) {
-                //$config_entry_type <: ensure_import_from(source = `.`),
+                $config_entry_type <: ensure_import_from(source = `.`),
             },
         },
 
@@ -63,13 +63,15 @@ multifile {
         },
         $domain_list += $domain,
     },
-    bubble($domain_list) file($name, $body) where {
+    bubble($domain_list) file($name, $body) as $file where {
+        $file <: before_each_file(),
         // migrate files
         $filename <: r".*components/([^/]+)/([^/]+)\.py$"($domain, $file_name),
         $domain_list <: includes $domain,
         $config_entry_type = capitalize($domain),
         $config_entry_type += "ConfigEntry",
         $body <: contains refactor_functions($config_entry_type, $file_name, $config_entry_type_defined),
+        $file <: after_each_file()
     },
 }
 ```
\ No newline at end of file

From f0387b6cb33fcf61ac2f8331e3dfc1922edfe38d Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Wed, 3 Jul 2024 08:19:30 +0000
Subject: [PATCH 5/5] Fix imports

---
 .grit/patterns/runtime_data_migration.md | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/.grit/patterns/runtime_data_migration.md b/.grit/patterns/runtime_data_migration.md
index 0c662a66fea..8824d37127c 100644
--- a/.grit/patterns/runtime_data_migration.md
+++ b/.grit/patterns/runtime_data_migration.md
@@ -8,6 +8,19 @@ Migrate an integration from hass.data to entry.runtime_data
 ```grit
 language python
 
+pattern import_from_custom($source) {
+    $name where {
+        $program <: module($statements),
+        $statements <: contains or {
+                import_from_statement(),
+                import_statement(),
+            } as $import,
+        // Ruff will handle double imports and sorting for us
+        // So we add always the import
+        $import => `from $source import $name\n$import`,
+    }
+}
+
 pattern refactor_functions($config_entry_type, $file_name, $config_entry_type_defined) {
     function_definition($parameters, $body) as $func where {
         // change config entry type
@@ -15,7 +28,7 @@ pattern refactor_functions($config_entry_type, $file_name, $config_entry_type_de
             $type <: type(type="ConfigEntry"),
             $type => $config_entry_type,
             if ($file_name <: not `__init__`) {
-                $config_entry_type <: ensure_import_from(source = `.`),
+                $config_entry_type <: import_from_custom(source = `.`),
             },
         },
 
@@ -63,15 +76,13 @@ multifile {
         },
         $domain_list += $domain,
     },
-    bubble($domain_list) file($name, $body) as $file where {
-        $file <: before_each_file(),
+    bubble($domain_list) file($name, $body) where {
         // migrate files
         $filename <: r".*components/([^/]+)/([^/]+)\.py$"($domain, $file_name),
         $domain_list <: includes $domain,
         $config_entry_type = capitalize($domain),
         $config_entry_type += "ConfigEntry",
         $body <: contains refactor_functions($config_entry_type, $file_name, $config_entry_type_defined),
-        $file <: after_each_file()
     },
 }
 ```
\ No newline at end of file