fix(python): regex injection not working for concatenated strings (#8197)

Co-authored-by: Riley Bruins <ribru17@hotmail.com>

7 files changed, 83 insertions, 3 deletions

diff --git a/runtime/queries/python/highlights.scm b/runtime/queries/python/highlights.scm
index 70f6d9acc..350990b1f 100644
--- a/runtime/queries/python/highlights.scm
+++ b/runtime/queries/python/highlights.scm
@@ -417,11 +417,19 @@
   function: (attribute
     object: (identifier) @_re)
   arguments: (argument_list
-    .
     (string
       (string_content) @string.regexp))
   (#eq? @_re "re"))
 
+(call
+  function: (attribute
+    object: (identifier) @_re)
+  arguments: (argument_list
+    (concatenated_string
+      (string
+        (string_content) @string.regexp)))
+  (#eq? @_re "re"))
+
 ; Decorators
 ((decorator
   "@" @attribute)
diff --git a/runtime/queries/python/injections.scm b/runtime/queries/python/injections.scm
index bbc924be4..2a120d9b5 100644
--- a/runtime/queries/python/injections.scm
+++ b/runtime/queries/python/injections.scm
@@ -2,12 +2,24 @@
   function: (attribute
     object: (identifier) @_re)
   arguments: (argument_list
-    .
     (string
       (string_content) @injection.content))
   (#eq? @_re "re")
   (#set! injection.language "regex"))
 
+(call
+  function: (attribute
+    object: (identifier) @_re)
+  arguments: (argument_list
+    (concatenated_string
+      [
+        (string
+          (string_content) @injection.content)
+        (comment)
+      ]+))
+  (#eq? @_re "re")
+  (#set! injection.language "regex"))
+
 ((binary_operator
   left: (string
     (string_content) @injection.content)
diff --git a/runtime/queries/starlark/injections.scm b/runtime/queries/starlark/injections.scm
index 0b920cbf9..a6656db97 100644
--- a/runtime/queries/starlark/injections.scm
+++ b/runtime/queries/starlark/injections.scm
@@ -1 +1,8 @@
-; inherits: python
+((binary_operator
+  left: (string
+    (string_content) @injection.content)
+  operator: "%")
+  (#set! injection.language "printf"))
+
+((comment) @injection.content
+  (#set! injection.language "comment"))
diff --git a/tests/query/highlights/python/regex.py b/tests/query/highlights/python/regex.py
new file mode 100644
index 000000000..385ec7760
--- /dev/null
+++ b/tests/query/highlights/python/regex.py
@@ -0,0 +1,19 @@
+import re
+
+re_test = re.compile(r"^(?P<year>\d{4}) (?P<day>\d) \w\s{,3}$")
+#                                    ^ @string.regexp
+re_test = re.compile(
+    # comment
+    # ^ @comment
+    r"^(?P<year>\d{4}){1}"
+    # ^ @string.regexp
+    # comment
+    # ^ @comment
+    r"(?P<day>\d) \w\s{,3}"
+    # ^ @string.regexp
+    # comment
+    # ^ @comment
+)
+# interpolation
+print("foo %s bar %d" % ("arg1", 2))
+#          ^ @character
diff --git a/tests/query/highlights/starlark/test.bzl b/tests/query/highlights/starlark/test.bzl
new file mode 100644
index 000000000..88251f615
--- /dev/null
+++ b/tests/query/highlights/starlark/test.bzl
@@ -0,0 +1,11 @@
+# split
+# ^ @comment
+assert_eq('foo bar'.split(' '), ['foo', 'bar'])
+# ^ @keyword
+#        ^ @punctuation.bracket
+#          ^ @string
+#                  ^ @punctuation.delimiter
+assert_eq('foo bar foo'.rsplit(' ', 1), ['foo bar', 'foo'])
+#                                   ^ @number
+assert_eq("foo %s bar %d" % ("arg1", 2), "foo arg1 bar 2")
+#              ^ @character
diff --git a/tests/query/injections/python/test.py b/tests/query/injections/python/test.py
new file mode 100644
index 000000000..1d6a0a465
--- /dev/null
+++ b/tests/query/injections/python/test.py
@@ -0,0 +1,18 @@
+import re
+
+re_test = re.compile(r"^(?P<year>\d{4}) (?P<day>\d) \w\s{,3}$")
+#                                    ^ @regex
+re_test = re.compile(
+    # comment
+    # ^ @comment
+    r"^(?P<year>\d{4}) "
+    # comment
+    # ^ @comment
+    r"(?P<day>\d) \w\s{,3}$"
+    # ^ @regex
+    # comment
+    # comment
+    # ^ @comment
+)
+print("foo %s bar %d" % ("arg1", 2))
+#          ^ @printf
diff --git a/tests/query/injections/starlark/test.bzl b/tests/query/injections/starlark/test.bzl
new file mode 100644
index 000000000..4a8af5172
--- /dev/null
+++ b/tests/query/injections/starlark/test.bzl
@@ -0,0 +1,5 @@
+# split
+# ^ @comment
+assert_eq('foo bar'.split(' '), ['foo', 'bar'])
+assert_eq("foo %s bar %d" % ("arg1", 2), "foo arg1 bar 2")
+#              ^ @printf