2 files changed, 36 insertions, 6 deletions
diff --git a/runtime/queries/python/folds.scm b/runtime/queries/python/folds.scm
index 7c547db38f..ecb9352d78 100644
--- a/runtime/queries/python/folds.scm
+++ b/runtime/queries/python/folds.scm
@@ -21,3 +21,8 @@
   (dictionary)
   (string)
 ] @fold
+
+[
+  (import_statement)
+  (import_from_statement)
+]+ @fold
diff --git a/runtime/queries/python/highlights.scm b/runtime/queries/python/highlights.scm
index 4d62aeff06..5e5a2a88de 100644
--- a/runtime/queries/python/highlights.scm
+++ b/runtime/queries/python/highlights.scm
@@ -188,7 +188,7 @@
 
 ((module
   .
-  (comment) @keyword.directive)
+  (comment) @keyword.directive @nospell)
   (#lua-match? @keyword.directive "^#!/"))
 
 (string) @string
@@ -204,19 +204,41 @@
   (comment)*
   .
   (expression_statement
-    (string) @string.documentation @spell))
+    (string) @string.documentation))
 
 (class_definition
   body: (block
     .
     (expression_statement
-      (string) @string.documentation @spell)))
+      (string) @string.documentation)))
 
 (function_definition
   body: (block
     .
     (expression_statement
-      (string) @string.documentation @spell)))
+      (string) @string.documentation)))
+
+(module
+  .
+  (comment)*
+  .
+  (expression_statement
+    (string
+      (string_content) @spell)))
+
+(class_definition
+  body: (block
+    .
+    (expression_statement
+      (string
+        (string_content) @spell))))
+
+(function_definition
+  body: (block
+    .
+    (expression_statement
+      (string
+        (string_content) @spell))))
 
 ; Tokens
 [
@@ -278,7 +300,6 @@
 
 [
   "assert"
-  "class"
   "exec"
   "global"
   "nonlocal"
@@ -286,10 +307,14 @@
   "print"
   "with"
   "as"
-  "type"
 ] @keyword
 
 [
+  "type"
+  "class"
+] @keyword.type
+
+[
   "async"
   "await"
 ] @keyword.coroutine