diff --git a/scripts/rewrite_imports.py b/scripts/rewrite_imports.py
index 71f83d8b946a..0ed8a310d0ef 100644
--- a/scripts/rewrite_imports.py
+++ b/scripts/rewrite_imports.py
@@ -49,8 +49,10 @@ def transform_old_to_new(line, old_module, new_module,
         "from {old_module} import ..."
     then checks if the line contains
         "import {old_module} ..."
-    and finally checks if the line starts with (ignoring whitespace)
+    then checks if the line starts with (ignoring whitespace)
         "{old_module} ..."
+    and finally checks if the line contians
+        "'some-dict-key': {old_module} ..."
 
     In any of these cases, "{old_module}" is replaced with "{new_module}".
     If none match, nothing is returned.
@@ -96,6 +98,11 @@ def transform_old_to_new(line, old_module, new_module,
         # Only replace the first instance of the old_module.
         return line.replace(old_module, new_module, 1)
 
+    # Finally check for usage in dictionaries.
+    if ': ' + old_module in line:
+        # Only replace the first instance of the old_module.
+        return line.replace(': ' + old_module, ': ' + new_module, 1)
+
 
 def transform_line(line):
     """Transforms an import line in a PB2 module.