From f452b01fb7e05b0d8e54a2d01c23e75ae998f90f Mon Sep 17 00:00:00 2001
From: Alex Povel <git@alexpovel.de>
Date: Sat, 2 Dec 2023 23:23:33 +0100
Subject: [PATCH] feat: Python strings

---
 src/scoping/langs/python.rs       | 13 +++++++++++++
 tests/langs/python/in/strings.py  | 11 +++++++++++
 tests/langs/python/mod.rs         |  1 +
 tests/langs/python/out/strings.py | 11 +++++++++++
 4 files changed, 36 insertions(+)
 create mode 100644 tests/langs/python/in/strings.py
 create mode 100644 tests/langs/python/out/strings.py

diff --git a/src/scoping/langs/python.rs b/src/scoping/langs/python.rs
index 0d5f0104..1a09cb47 100644
--- a/src/scoping/langs/python.rs
+++ b/src/scoping/langs/python.rs
@@ -14,6 +14,8 @@ pub type PythonQuery = CodeQuery<CustomPythonQuery, PremadePythonQuery>;
 pub enum PremadePythonQuery {
     /// Comments.
     Comments,
+    /// Strings (raw, byte, f-strings; interpolation is respected; quotes included).
+    Strings,
     /// Docstrings (not including multi-line strings).
     DocStrings,
     /// Function names, at the definition site.
@@ -28,6 +30,17 @@ impl From<PremadePythonQuery> for TSQuery {
             Python::lang(),
             match value {
                 PremadePythonQuery::Comments => "(comment) @comment",
+                PremadePythonQuery::Strings => {
+                    // Match either normal `string`s or `string`s with `interpolation`;
+                    // using only the latter doesn't include the former.
+                    r#"
+                    [
+                        (string)
+                        (string (interpolation) @IGNORE)
+                    ]
+                    @string
+                    "#
+                }
                 PremadePythonQuery::DocStrings => {
                     // Triple-quotes are also used for multi-line strings. So look only
                     // for stand-alone expressions, which are not part of some variable
diff --git a/tests/langs/python/in/strings.py b/tests/langs/python/in/strings.py
new file mode 100644
index 00000000..e5f13224
--- /dev/null
+++ b/tests/langs/python/in/strings.py
@@ -0,0 +1,11 @@
+name = "Ali__T__ce"
+ag__T__e = 30
+f_string = f"name:__T__ {name}, age: {ag__T__e}__T__"
+
+multiline_f_string = f"""This is a__T__
+multiline{f_string} string
+spanning several lines__T__"""
+
+raw_string = r"This is a raw string__T__ with no special treatment for \n"
+bytes_string = b"This is a bytes__T__ string"
+bytes_string = rf"This is a __T__raw f-string with {raw_string}"
diff --git a/tests/langs/python/mod.rs b/tests/langs/python/mod.rs
index 4150d249..72f35b7b 100644
--- a/tests/langs/python/mod.rs
+++ b/tests/langs/python/mod.rs
@@ -5,6 +5,7 @@ use super::{get_input_output, nuke_target};
 
 #[rstest]
 #[case("docstring.py", PythonQuery::Premade(PremadePythonQuery::DocStrings))]
+#[case("strings.py", PythonQuery::Premade(PremadePythonQuery::Strings))]
 #[case("comments-lf.py", PythonQuery::Premade(PremadePythonQuery::Comments))]
 #[case("comments-crlf.py", PythonQuery::Premade(PremadePythonQuery::Comments))]
 #[case(
diff --git a/tests/langs/python/out/strings.py b/tests/langs/python/out/strings.py
new file mode 100644
index 00000000..196eb77b
--- /dev/null
+++ b/tests/langs/python/out/strings.py
@@ -0,0 +1,11 @@
+name = "Alice"
+ag__T__e = 30
+f_string = f"name: {name}, age: {ag__T__e}"
+
+multiline_f_string = f"""This is a
+multiline{f_string} string
+spanning several lines"""
+
+raw_string = r"This is a raw string with no special treatment for \n"
+bytes_string = b"This is a bytes string"
+bytes_string = rf"This is a raw f-string with {raw_string}"