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