ruff_python_ast/
script.rs

1use std::sync::LazyLock;
2
3use memchr::memmem::Finder;
4use ruff_source_file::UniversalNewlineIterator;
5use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
6
7static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(b"# /// script"));
8
9/// PEP 723 metadata as parsed from a `script` comment block.
10///
11/// See: <https://peps.python.org/pep-0723/>
12///
13/// Vendored from: <https://github.com/astral-sh/uv/blob/debe67ffdb0cd7835734100e909b2d8f79613743/crates/uv-scripts/src/lib.rs#L283>
14#[derive(Debug, Clone, Eq, PartialEq)]
15pub struct ScriptTag {
16    /// The metadata block.
17    metadata: String,
18    /// The source range of the metadata block, including its opening and closing delimiters.
19    range: TextRange,
20    /// Maps offsets in the extracted metadata to offsets in the original Python script.
21    source_map: ScriptSourceMap,
22}
23
24impl ScriptTag {
25    /// Returns the TOML contents of the metadata block.
26    pub fn metadata(&self) -> &str {
27        &self.metadata
28    }
29
30    /// Returns the map from extracted TOML offsets to their original script offsets.
31    pub fn source_map(&self) -> &ScriptSourceMap {
32        &self.source_map
33    }
34
35    /// Given the contents of a Python file, extract the `script` metadata block with leading
36    /// comment hashes removed and map its offsets to the original Python script.
37    ///
38    /// Given the following input string representing the contents of a Python script:
39    ///
40    /// ```python
41    /// #!/usr/bin/env python3
42    /// # /// script
43    /// # requires-python = '>=3.11'
44    /// # dependencies = [
45    /// #   'requests<3',
46    /// #   'rich',
47    /// # ]
48    /// # ///
49    ///
50    /// import requests
51    ///
52    /// print("Hello, World!")
53    /// ```
54    ///
55    /// This function extracts the metadata:
56    /// ```toml
57    /// requires-python = '>=3.11'
58    /// dependencies = [
59    ///   'requests<3',
60    ///   'rich',
61    /// ]
62    /// ```
63    ///
64    /// See: <https://peps.python.org/pep-0723/>
65    pub fn parse(contents: &[u8]) -> Option<Self> {
66        FINDER
67            .find_iter(contents)
68            .find_map(|index| Self::parse_at(contents, index))
69    }
70
71    fn parse_at(contents: &[u8], index: usize) -> Option<Self> {
72        // The opening pragma must be the first line, or immediately preceded by a newline.
73        if !(index == 0 || matches!(contents[index - 1], b'\r' | b'\n')) {
74            return None;
75        }
76
77        let contents = std::str::from_utf8(contents).ok()?;
78        let contents = &contents[index..];
79
80        let start = TextSize::try_from(index).ok()?;
81        let mut lines = UniversalNewlineIterator::with_offset(contents, start);
82
83        // Ensure that the first line is exactly `# /// script`.
84        if lines.next().is_none_or(|line| line != "# /// script") {
85            return None;
86        }
87
88        // > Every line between these two lines (# /// TYPE and # ///) MUST be a comment starting
89        // > with #. If there are characters after the # then the first character MUST be a space. The
90        // > embedded content is formed by taking away the first two characters of each line if the
91        // > second character is a space, otherwise just the first character (which means the line
92        // > consists of only a single #).
93        let mut metadata = String::new();
94        let mut source_map = ScriptSourceMap::default();
95        let mut closing = None;
96
97        for line in lines {
98            // Remove the leading `#`.
99            let Some(comment) = line.strip_prefix('#') else {
100                break;
101            };
102
103            let (content, indent_len) = if comment.is_empty() {
104                ("", TextSize::ZERO)
105            } else if let Some(content) = comment.strip_prefix(' ') {
106                (content, ' '.text_len())
107            } else {
108                break;
109            };
110
111            if content == "///" {
112                closing = Some((metadata.len(), source_map.markers.len(), line.range()));
113            }
114
115            let prefix_length = '#'.text_len() + indent_len;
116
117            source_map.push_marker(metadata.text_len(), line.start() + prefix_length);
118            metadata.push_str(content);
119            metadata.push('\n');
120        }
121
122        // The last closing `# ///` wins, so discard that delimiter and everything after it.
123        //
124        // For example, given:
125        // ```python
126        // # /// script
127        // #
128        // # ///
129        // #
130        // # ///
131        // ```
132        //
133        // The latter `///` is the closing pragma
134        let (metadata_end, marker_count, closing_range) = closing?;
135        metadata.truncate(metadata_end);
136        source_map.truncate(marker_count);
137
138        if metadata.is_empty() {
139            metadata.push('\n');
140        } else {
141            source_map.push_marker(metadata.text_len(), closing_range.start());
142        }
143
144        Some(Self {
145            metadata,
146            range: TextRange::new(start, closing_range.end()),
147            source_map,
148        })
149    }
150}
151
152impl Ranged for ScriptTag {
153    fn range(&self) -> TextRange {
154        self.range
155    }
156}
157
158/// Maps offsets in extracted script metadata to offsets in the original Python source.
159#[derive(Clone, Debug, Default, Eq, PartialEq)]
160pub struct ScriptSourceMap {
161    markers: Vec<ScriptSourceMarker>,
162}
163
164impl ScriptSourceMap {
165    /// Maps a metadata offset to the corresponding offset in the Python source.
166    pub fn map_offset(&self, offset: TextSize) -> TextSize {
167        let Some(index) = self
168            .markers
169            .partition_point(|marker| marker.metadata_offset <= offset)
170            .checked_sub(1)
171        else {
172            return offset;
173        };
174        let marker = &self.markers[index];
175
176        marker.source_offset + (offset - marker.metadata_offset)
177    }
178
179    /// Maps a metadata range to its corresponding range in the Python source.
180    pub fn map_range(&self, range: TextRange) -> TextRange {
181        TextRange::new(self.map_offset(range.start()), self.map_offset(range.end()))
182    }
183
184    fn push_marker(&mut self, metadata_offset: TextSize, source_offset: TextSize) {
185        self.markers.push(ScriptSourceMarker {
186            metadata_offset,
187            source_offset,
188        });
189    }
190
191    fn truncate(&mut self, len: usize) {
192        self.markers.truncate(len);
193    }
194}
195
196#[derive(Clone, Debug, Eq, PartialEq)]
197struct ScriptSourceMarker {
198    metadata_offset: TextSize,
199    source_offset: TextSize,
200}
201
202#[cfg(test)]
203mod tests {
204    use ruff_text_size::{Ranged, TextLen, TextRange};
205
206    use super::ScriptTag;
207
208    #[test]
209    fn carriage_return_line_endings() -> Result<(), &'static str> {
210        let tag = ScriptTag::parse(b"# /// script\r# value = true\r# ///\r")
211            .ok_or("Expected script metadata with carriage-return line endings")?;
212
213        assert_eq!(tag.metadata(), "value = true\n");
214
215        Ok(())
216    }
217
218    #[test]
219    fn metadata_block_range_includes_both_delimiters() -> Result<(), &'static str> {
220        let prefix = "#!/usr/bin/env python3\n\n";
221        let metadata = "# /// script\n# dependencies = []\n# ///";
222        let source = format!("{prefix}{metadata}\n\nprint('hello')\n");
223        let tag = ScriptTag::parse(source.as_bytes()).ok_or("Expected valid script metadata")?;
224
225        assert_eq!(
226            tag.range(),
227            TextRange::at(prefix.text_len(), metadata.text_len())
228        );
229
230        Ok(())
231    }
232
233    #[test]
234    fn metadata_range_accounts_for_unicode_crlf_and_multiline_values() -> Result<(), &'static str> {
235        let metadata_value = r#""""
236first
237
238last
239""""#;
240        let source_value = r#""""
241# first
242#
243# last
244# """"#
245            .replace('\n', "\r\n");
246        let source = format!("π\r\n# /// script\r\n# value = {source_value}\r\n# ///\r\n");
247        let tag = ScriptTag::parse(source.as_bytes()).ok_or("Expected valid script metadata")?;
248
249        assert_eq!(tag.metadata(), format!("value = {metadata_value}\n"));
250
251        let metadata_range = TextRange::at("value = ".text_len(), metadata_value.text_len());
252        let source_range = TextRange::at(
253            "π\r\n# /// script\r\n# value = ".text_len(),
254            source_value.text_len(),
255        );
256
257        assert_eq!(tag.source_map().map_range(metadata_range), source_range);
258
259        Ok(())
260    }
261
262    #[test]
263    fn last_closing_delimiter_discards_following_comments() -> Result<(), &'static str> {
264        let source = r"# /// script
265# first = true
266# ///
267# last = true
268# ///
269# ignored = true
270";
271        let tag = ScriptTag::parse(source.as_bytes()).ok_or("Expected valid script metadata")?;
272
273        assert_eq!(
274            tag.metadata(),
275            r"first = true
276///
277last = true
278"
279        );
280
281        let closing_start = source
282            .rfind("# ///")
283            .map(|offset| source[..offset].text_len())
284            .ok_or("Expected the final closing delimiter")?;
285        assert_eq!(
286            tag.source_map().map_offset(tag.metadata().text_len()),
287            closing_start,
288        );
289        assert_eq!(tag.end(), closing_start + "# ///".text_len());
290
291        Ok(())
292    }
293}
ruff_python_ast/script.rs

ruff_python_ast/
script.rs
