Skip to main content

uv_scripts/
lib.rs

1use std::borrow::Cow;
2use std::collections::BTreeMap;
3use std::io;
4use std::path::{Path, PathBuf};
5use std::str::FromStr;
6use std::sync::LazyLock;
7
8use memchr::memmem::Finder;
9use serde::Deserialize;
10use thiserror::Error;
11use tracing::instrument;
12use url::Url;
13
14use uv_configuration::NoSources;
15use uv_normalize::PackageName;
16use uv_pep440::VersionSpecifiers;
17use uv_pypi_types::VerbatimParsedUrl;
18use uv_redacted::DisplaySafeUrl;
19use uv_settings::{GlobalOptions, ResolverInstallerSchema};
20use uv_warnings::warn_user;
21use uv_workspace::pyproject::{ExtraBuildDependency, Sources};
22
23pub use uv_configuration::ExcludeDependency;
24pub use uv_workspace::pyproject::OverrideDependency;
25
26static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(b"# /// script"));
27
28/// A PEP 723 item, either read from a script on disk or provided via `stdin`.
29#[derive(Debug)]
30pub enum Pep723Item {
31    /// A PEP 723 script read from disk.
32    Script(Pep723Script),
33    /// A PEP 723 script provided via `stdin`.
34    Stdin(Pep723Metadata),
35    /// A PEP 723 script provided via a remote URL.
36    Remote(Pep723Metadata, DisplaySafeUrl),
37}
38
39impl Pep723Item {
40    /// Return the [`Pep723Metadata`] associated with the item.
41    pub fn metadata(&self) -> &Pep723Metadata {
42        match self {
43            Self::Script(script) => &script.metadata,
44            Self::Stdin(metadata) => metadata,
45            Self::Remote(metadata, ..) => metadata,
46        }
47    }
48
49    /// Return the PEP 723 script, if any.
50    pub fn as_script(&self) -> Option<&Pep723Script> {
51        match self {
52            Self::Script(script) => Some(script),
53            _ => None,
54        }
55    }
56}
57
58/// A reference to a PEP 723 item.
59#[derive(Debug, Copy, Clone)]
60pub enum Pep723ItemRef<'item> {
61    /// A PEP 723 script read from disk.
62    Script(&'item Pep723Script),
63    /// A PEP 723 script provided via `stdin`.
64    Stdin(&'item Pep723Metadata),
65    /// A PEP 723 script provided via a remote URL.
66    Remote(&'item Pep723Metadata, &'item Url),
67}
68
69impl Pep723ItemRef<'_> {
70    /// Return the [`Pep723Metadata`] associated with the item.
71    pub fn metadata(&self) -> &Pep723Metadata {
72        match self {
73            Self::Script(script) => &script.metadata,
74            Self::Stdin(metadata) => metadata,
75            Self::Remote(metadata, ..) => metadata,
76        }
77    }
78
79    /// Return the path of the PEP 723 item, if any.
80    pub fn path(&self) -> Option<&Path> {
81        match self {
82            Self::Script(script) => Some(&script.path),
83            Self::Stdin(..) => None,
84            Self::Remote(..) => None,
85        }
86    }
87
88    /// Determine the working directory for the script.
89    pub fn directory(&self) -> Result<PathBuf, io::Error> {
90        match self {
91            Self::Script(script) => Ok(std::path::absolute(&script.path)?
92                .parent()
93                .expect("script path has no parent")
94                .to_owned()),
95            Self::Stdin(..) | Self::Remote(..) => std::env::current_dir(),
96        }
97    }
98
99    /// Collect any `tool.uv.index` from the script.
100    pub fn indexes(&self, source_strategy: &NoSources) -> &[uv_distribution_types::Index] {
101        match source_strategy {
102            NoSources::None | NoSources::Packages(_) => self
103                .metadata()
104                .tool
105                .as_ref()
106                .and_then(|tool| tool.uv.as_ref())
107                .and_then(|uv| uv.top_level.index.as_deref())
108                .unwrap_or(&[]),
109            NoSources::All => &[],
110        }
111    }
112
113    /// Collect any `tool.uv.sources` from the script.
114    pub fn sources(&self, source_strategy: &NoSources) -> Cow<'_, BTreeMap<PackageName, Sources>> {
115        static EMPTY: BTreeMap<PackageName, Sources> = BTreeMap::new();
116        let sources = self
117            .metadata()
118            .tool
119            .as_ref()
120            .and_then(|tool| tool.uv.as_ref())
121            .and_then(|uv| uv.sources.as_ref())
122            .unwrap_or(&EMPTY);
123
124        match source_strategy {
125            NoSources::None => Cow::Borrowed(sources),
126            NoSources::All => Cow::Borrowed(&EMPTY),
127            NoSources::Packages(packages) => Cow::Owned(
128                sources
129                    .iter()
130                    .filter(|(name, _)| !packages.contains(name))
131                    .map(|(name, sources)| (name.clone(), sources.clone()))
132                    .collect(),
133            ),
134        }
135    }
136}
137
138impl<'item> From<&'item Pep723Item> for Pep723ItemRef<'item> {
139    fn from(item: &'item Pep723Item) -> Self {
140        match item {
141            Pep723Item::Script(script) => Self::Script(script),
142            Pep723Item::Stdin(metadata) => Self::Stdin(metadata),
143            Pep723Item::Remote(metadata, url) => Self::Remote(metadata, url),
144        }
145    }
146}
147
148impl<'item> From<&'item Pep723Script> for Pep723ItemRef<'item> {
149    fn from(script: &'item Pep723Script) -> Self {
150        Self::Script(script)
151    }
152}
153
154/// A PEP 723 script, including its [`Pep723Metadata`].
155#[derive(Debug, Clone)]
156pub struct Pep723Script {
157    /// The path to the Python script.
158    pub path: PathBuf,
159    /// The parsed [`Pep723Metadata`] table from the script.
160    pub metadata: Pep723Metadata,
161    /// The content of the script before the metadata table.
162    pub prelude: String,
163    /// The content of the script after the metadata table.
164    pub postlude: String,
165}
166
167impl Pep723Script {
168    /// Read the PEP 723 `script` metadata from a Python file, if it exists.
169    ///
170    /// Returns `None` if the file is missing a PEP 723 metadata block.
171    ///
172    /// See: <https://peps.python.org/pep-0723/>
173    pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
174        let contents = fs_err::tokio::read(&file).await?;
175
176        // Extract the `script` tag.
177        let ScriptTag {
178            prelude,
179            metadata,
180            postlude,
181        } = match ScriptTag::parse(&contents) {
182            Ok(Some(tag)) => tag,
183            Ok(None) => return Ok(None),
184            Err(err) => return Err(err),
185        };
186
187        // Parse the metadata.
188        let metadata = Pep723Metadata::from_str(&metadata)?;
189
190        Ok(Some(Self {
191            path: std::path::absolute(file)?,
192            metadata,
193            prelude,
194            postlude,
195        }))
196    }
197
198    /// Reads a Python script and generates a default PEP 723 metadata table.
199    ///
200    /// See: <https://peps.python.org/pep-0723/>
201    pub async fn init(
202        file: impl AsRef<Path>,
203        requires_python: &VersionSpecifiers,
204    ) -> Result<Self, Pep723Error> {
205        let contents = fs_err::tokio::read(&file).await?;
206        let (prelude, metadata, postlude) = Self::init_metadata(&contents, requires_python)?;
207        Ok(Self {
208            path: std::path::absolute(file)?,
209            metadata,
210            prelude,
211            postlude,
212        })
213    }
214
215    /// Generates a default PEP 723 metadata table from the provided script contents.
216    ///
217    /// See: <https://peps.python.org/pep-0723/>
218    fn init_metadata(
219        contents: &[u8],
220        requires_python: &VersionSpecifiers,
221    ) -> Result<(String, Pep723Metadata, String), Pep723Error> {
222        // Define the default metadata.
223        let default_metadata = if requires_python.is_empty() {
224            indoc::formatdoc! {r"
225                dependencies = []
226            ",
227            }
228        } else {
229            indoc::formatdoc! {r#"
230                requires-python = "{requires_python}"
231                dependencies = []
232                "#,
233                requires_python = requires_python,
234            }
235        };
236        let metadata = Pep723Metadata::from_str(&default_metadata)?;
237
238        // Extract the shebang and script content.
239        let (shebang, postlude) = extract_shebang(contents)?;
240
241        // Add a newline to the beginning if it starts with a valid metadata comment line.
242        let postlude = if postlude.strip_prefix('#').is_some_and(|postlude| {
243            postlude
244                .chars()
245                .next()
246                .is_some_and(|c| matches!(c, ' ' | '\r' | '\n'))
247        }) {
248            format!("\n{postlude}")
249        } else {
250            postlude
251        };
252
253        Ok((
254            if shebang.is_empty() {
255                String::new()
256            } else {
257                format!("{shebang}\n")
258            },
259            metadata,
260            postlude,
261        ))
262    }
263
264    /// Create a PEP 723 script at the given path.
265    pub async fn create(
266        file: impl AsRef<Path>,
267        requires_python: &VersionSpecifiers,
268        existing_contents: Option<Vec<u8>>,
269        bare: bool,
270    ) -> Result<(), Pep723Error> {
271        let file = file.as_ref();
272
273        let script_name = file
274            .file_name()
275            .and_then(|name| name.to_str())
276            .ok_or_else(|| Pep723Error::InvalidFilename(file.to_string_lossy().to_string()))?;
277
278        let default_metadata = indoc::formatdoc! {r#"
279            requires-python = "{requires_python}"
280            dependencies = []
281            "#,
282        };
283        let metadata = serialize_metadata(&default_metadata);
284
285        let script = if let Some(existing_contents) = existing_contents {
286            let (mut shebang, contents) = extract_shebang(&existing_contents)?;
287            if !shebang.is_empty() {
288                shebang.push_str("\n#\n");
289                // If the shebang doesn't contain `uv`, it's probably something like
290                // `#! /usr/bin/env python`, which isn't going to respect the inline metadata.
291                // Issue a warning for users who might not know that.
292                // TODO: There are a lot of mistakes we could consider detecting here, like
293                // `uv run` without `--script` when the file doesn't end in `.py`.
294                if !regex::regex!(r"\buv\b").is_match(&shebang) {
295                    warn_user!(
296                        "If you execute {} directly, it might ignore its inline metadata.\nConsider replacing its shebang with: {}",
297                        file.to_string_lossy().cyan(),
298                        "#!/usr/bin/env -S uv run --script".cyan(),
299                    );
300                }
301            }
302            indoc::formatdoc! {r"
303            {shebang}{metadata}
304            {contents}" }
305        } else if bare {
306            metadata
307        } else {
308            indoc::formatdoc! {r#"
309            {metadata}
310
311            def main() -> None:
312                print("Hello from {name}!")
313
314
315            if __name__ == "__main__":
316                main()
317        "#,
318                metadata = metadata,
319                name = script_name,
320            }
321        };
322
323        Ok(fs_err::tokio::write(file, script).await?)
324    }
325
326    /// Replace the existing metadata in the file with new metadata and write the updated content.
327    pub fn write(&self, metadata: &str) -> Result<(), io::Error> {
328        let content = format!(
329            "{}{}{}",
330            self.prelude,
331            serialize_metadata(metadata),
332            self.postlude
333        );
334
335        fs_err::write(&self.path, content)?;
336
337        Ok(())
338    }
339
340    /// Return the [`Sources`] defined in the PEP 723 metadata.
341    pub fn sources(&self) -> &BTreeMap<PackageName, Sources> {
342        static EMPTY: BTreeMap<PackageName, Sources> = BTreeMap::new();
343
344        self.metadata
345            .tool
346            .as_ref()
347            .and_then(|tool| tool.uv.as_ref())
348            .and_then(|uv| uv.sources.as_ref())
349            .unwrap_or(&EMPTY)
350    }
351}
352
353/// PEP 723 metadata as parsed from a `script` comment block.
354///
355/// See: <https://peps.python.org/pep-0723/>
356#[derive(Debug, Deserialize, Clone)]
357#[serde(rename_all = "kebab-case")]
358pub struct Pep723Metadata {
359    pub dependencies: Option<Vec<uv_pep508::Requirement<VerbatimParsedUrl>>>,
360    pub requires_python: Option<VersionSpecifiers>,
361    pub tool: Option<Tool>,
362    /// The raw unserialized document.
363    #[serde(skip)]
364    pub raw: String,
365}
366
367impl Pep723Metadata {
368    /// Parse the PEP 723 metadata from `stdin`.
369    pub fn parse(contents: &[u8]) -> Result<Option<Self>, Pep723Error> {
370        // Extract the `script` tag.
371        let ScriptTag { metadata, .. } = match ScriptTag::parse(contents) {
372            Ok(Some(tag)) => tag,
373            Ok(None) => return Ok(None),
374            Err(err) => return Err(err),
375        };
376
377        // Parse the metadata.
378        Ok(Some(Self::from_str(&metadata)?))
379    }
380
381    /// Read the PEP 723 `script` metadata from a Python file, if it exists.
382    ///
383    /// Returns `None` if the file is missing a PEP 723 metadata block.
384    ///
385    /// See: <https://peps.python.org/pep-0723/>
386    pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
387        let contents = fs_err::tokio::read(&file).await?;
388
389        // Extract the `script` tag.
390        let ScriptTag { metadata, .. } = match ScriptTag::parse(&contents) {
391            Ok(Some(tag)) => tag,
392            Ok(None) => return Ok(None),
393            Err(err) => return Err(err),
394        };
395
396        // Parse the metadata.
397        Ok(Some(Self::from_str(&metadata)?))
398    }
399}
400
401impl FromStr for Pep723Metadata {
402    type Err = toml::de::Error;
403
404    /// Parse `Pep723Metadata` from a raw TOML string.
405    #[instrument(name = "toml::from_str PEP 723 metadata", skip_all)]
406    fn from_str(raw: &str) -> Result<Self, Self::Err> {
407        let metadata = toml::from_str(raw)?;
408        Ok(Self {
409            raw: raw.to_string(),
410            ..metadata
411        })
412    }
413}
414
415#[derive(Deserialize, Debug, Clone)]
416#[serde(rename_all = "kebab-case")]
417pub struct Tool {
418    pub uv: Option<ToolUv>,
419}
420
421#[derive(Debug, Deserialize, Clone)]
422#[serde(deny_unknown_fields, rename_all = "kebab-case")]
423pub struct ToolUv {
424    #[serde(flatten)]
425    pub globals: GlobalOptions,
426    #[serde(flatten)]
427    pub top_level: ResolverInstallerSchema,
428    pub override_dependencies: Option<Vec<OverrideDependency>>,
429    pub exclude_dependencies: Option<Vec<ExcludeDependency>>,
430    pub constraint_dependencies: Option<Vec<uv_pep508::Requirement<VerbatimParsedUrl>>>,
431    pub build_constraint_dependencies: Option<Vec<uv_pep508::Requirement<VerbatimParsedUrl>>>,
432    pub extra_build_dependencies: Option<BTreeMap<PackageName, Vec<ExtraBuildDependency>>>,
433    pub sources: Option<BTreeMap<PackageName, Sources>>,
434}
435
436#[derive(Debug, Error)]
437pub enum Pep723Error {
438    #[error(
439        "An opening tag (`# /// script`) was found without a closing tag (`# ///`). Ensure that every line between the opening and closing tags (including empty lines) starts with a leading `#`."
440    )]
441    UnclosedBlock,
442    #[error(
443        "An opening tag (`# /// script`) was found, but the closing tag (`# ///`) has trailing content. Remove the trailing content so the line is exactly `# ///`."
444    )]
445    UnclosedBlockTrailingContent,
446    #[error("The script contains multiple PEP 723 metadata blocks")]
447    DuplicateBlock,
448    #[error("The PEP 723 metadata block is missing from the script.")]
449    MissingTag,
450    #[error(transparent)]
451    Io(#[from] io::Error),
452    #[error(transparent)]
453    Utf8(#[from] std::str::Utf8Error),
454    #[error(transparent)]
455    Toml(#[from] toml::de::Error),
456    #[error("Invalid filename `{0}` supplied")]
457    InvalidFilename(String),
458}
459
460#[derive(Debug, Clone, Eq, PartialEq)]
461pub struct ScriptTag {
462    /// The content of the script before the metadata block.
463    prelude: String,
464    /// The metadata block.
465    metadata: String,
466    /// The content of the script after the metadata block.
467    postlude: String,
468}
469
470impl ScriptTag {
471    /// Given the contents of a Python file, extract the `script` metadata block with leading
472    /// comment hashes removed, any preceding shebang or content (prelude), and the remaining Python
473    /// script.
474    ///
475    /// Given the following input string representing the contents of a Python script:
476    ///
477    /// ```python
478    /// #!/usr/bin/env python3
479    /// # /// script
480    /// # requires-python = '>=3.11'
481    /// # dependencies = [
482    /// #   'requests<3',
483    /// #   'rich',
484    /// # ]
485    /// # ///
486    ///
487    /// import requests
488    ///
489    /// print("Hello, World!")
490    /// ```
491    ///
492    /// This function would return:
493    ///
494    /// - Preamble: `#!/usr/bin/env python3\n`
495    /// - Metadata: `requires-python = '>=3.11'\ndependencies = [\n  'requests<3',\n  'rich',\n]`
496    /// - Postlude: `import requests\n\nprint("Hello, World!")\n`
497    ///
498    /// See: <https://peps.python.org/pep-0723/>
499    pub fn parse(contents: &[u8]) -> Result<Option<Self>, Pep723Error> {
500        // Identify the opening pragma.
501        let Some(index) = FINDER.find(contents) else {
502            return Ok(None);
503        };
504
505        // The opening pragma must be the first line, or immediately preceded by a newline.
506        if !(index == 0 || matches!(contents[index - 1], b'\r' | b'\n')) {
507            return Ok(None);
508        }
509
510        // Extract the preceding content.
511        let prelude = std::str::from_utf8(&contents[..index])?;
512
513        // Decode as UTF-8.
514        let contents = &contents[index..];
515        let contents = std::str::from_utf8(contents)?;
516
517        let mut lines = contents.lines();
518
519        // Ensure that the first line is exactly `# /// script`.
520        if lines.next().is_none_or(|line| line != "# /// script") {
521            return Ok(None);
522        }
523
524        // > Every line between these two lines (# /// TYPE and # ///) MUST be a comment starting
525        // > with #. If there are characters after the # then the first character MUST be a space. The
526        // > embedded content is formed by taking away the first two characters of each line if the
527        // > second character is a space, otherwise just the first character (which means the line
528        // > consists of only a single #).
529        let mut toml = vec![];
530
531        for line in lines {
532            // Remove the leading `#`.
533            let Some(line) = line.strip_prefix('#') else {
534                break;
535            };
536
537            // If the line is empty, continue.
538            if line.is_empty() {
539                toml.push("");
540                continue;
541            }
542
543            // Otherwise, the line _must_ start with ` `.
544            let Some(line) = line.strip_prefix(' ') else {
545                break;
546            };
547
548            toml.push(line);
549        }
550
551        // Find the closing `# ///`. The precedence is such that we need to identify the _last_ such
552        // line.
553        //
554        // For example, given:
555        // ```python
556        // # /// script
557        // #
558        // # ///
559        // #
560        // # ///
561        // ```
562        //
563        // The latter `///` is the closing pragma. Track malformed terminators while searching, but
564        // continue looking for an exact terminator so trailing comments cannot invalidate it.
565        let mut has_trailing_content = false;
566        let mut closing_index = None;
567
568        for (index, line) in toml.iter().enumerate().rev() {
569            if *line == "///" {
570                closing_index = Some(index + 1);
571                break;
572            }
573
574            if line.starts_with("///") {
575                has_trailing_content = true;
576            }
577        }
578
579        let Some(index) = closing_index else {
580            return Err(if has_trailing_content {
581                Pep723Error::UnclosedBlockTrailingContent
582            } else {
583                Pep723Error::UnclosedBlock
584            });
585        };
586
587        // Discard any lines after the closing `# ///`.
588        //
589        // For example, given:
590        // ```python
591        // # /// script
592        // #
593        // # ///
594        // #
595        // #
596        // ```
597        //
598        // We need to discard the last two lines.
599        toml.truncate(index - 1);
600
601        // Extract the remaining content.
602        let postlude = contents.lines().skip(index + 1).collect::<Vec<_>>();
603
604        // Ensure that the remaining content doesn't include another complete `script` block.
605        // A `# /// script` line can be embedded content inside another typed block.
606        let mut lines = postlude.iter().peekable();
607        while let Some(line) = lines.next() {
608            // Capture the metadata.
609            let Some(metadata_type) = line.strip_prefix("# /// ") else {
610                continue;
611            };
612
613            // Parse the metadata type per spec
614            if metadata_type.is_empty()
615                || !metadata_type
616                    .bytes()
617                    .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-')
618            {
619                continue;
620            }
621
622            let is_script_block = metadata_type == "script";
623            let mut is_closed = false;
624            while let Some(line) = lines.next() {
625                // Per e.g. # dependencies = []
626                let Some(content) = line.strip_prefix('#') else {
627                    break;
628                };
629                if !(content.is_empty() || content.starts_with(' ')) {
630                    break;
631                }
632
633                if *line == "# ///" {
634                    let Some(next_line) = lines.peek() else {
635                        is_closed = true;
636                        break;
637                    };
638
639                    let Some(next_content) = next_line.strip_prefix('#') else {
640                        is_closed = true;
641                        break;
642                    };
643
644                    if !(next_content.is_empty() || next_content.starts_with(' ')) {
645                        is_closed = true;
646                        break;
647                    }
648                }
649            }
650
651            if is_script_block && is_closed {
652                return Err(Pep723Error::DuplicateBlock);
653            }
654        }
655
656        // Join the lines into a single string.
657        let prelude = prelude.to_string();
658        let metadata = toml.join("\n") + "\n";
659        let postlude = postlude.join("\n") + "\n";
660
661        Ok(Some(Self {
662            prelude,
663            metadata,
664            postlude,
665        }))
666    }
667}
668
669/// Extracts the shebang line from the given file contents and returns it along with the remaining
670/// content.
671fn extract_shebang(contents: &[u8]) -> Result<(String, String), Pep723Error> {
672    let contents = std::str::from_utf8(contents)?;
673
674    if contents.starts_with("#!") {
675        // Find the first newline.
676        let bytes = contents.as_bytes();
677        let index = bytes
678            .iter()
679            .position(|&b| b == b'\r' || b == b'\n')
680            .unwrap_or(bytes.len());
681
682        // Support `\r`, `\n`, and `\r\n` line endings.
683        let width = match bytes.get(index) {
684            Some(b'\r') => {
685                if bytes.get(index + 1) == Some(&b'\n') {
686                    2
687                } else {
688                    1
689                }
690            }
691            Some(b'\n') => 1,
692            _ => 0,
693        };
694
695        // Extract the shebang line.
696        let shebang = contents[..index].to_string();
697        let script = contents[index + width..].to_string();
698
699        Ok((shebang, script))
700    } else {
701        Ok((String::new(), contents.to_string()))
702    }
703}
704
705/// Formats the provided metadata by prefixing each line with `#` and wrapping it with script markers.
706fn serialize_metadata(metadata: &str) -> String {
707    let mut output = String::with_capacity(metadata.len() + 32);
708
709    output.push_str("# /// script");
710    output.push('\n');
711
712    for line in metadata.lines() {
713        output.push('#');
714        if !line.is_empty() {
715            output.push(' ');
716            output.push_str(line);
717        }
718        output.push('\n');
719    }
720
721    output.push_str("# ///");
722    output.push('\n');
723
724    output
725}
726
727#[cfg(test)]
728mod tests {
729    use std::assert_matches;
730
731    use crate::{Pep723Error, Pep723Script, ScriptTag, serialize_metadata};
732    use std::str::FromStr;
733
734    #[test]
735    fn missing_space() {
736        let contents = indoc::indoc! {r"
737        # /// script
738        #requires-python = '>=3.11'
739        # ///
740    "};
741
742        assert_matches!(
743            ScriptTag::parse(contents.as_bytes()),
744            Err(Pep723Error::UnclosedBlock)
745        );
746    }
747
748    #[test]
749    fn no_closing_pragma() {
750        let contents = indoc::indoc! {r"
751        # /// script
752        # requires-python = '>=3.11'
753        # dependencies = [
754        #     'requests<3',
755        #     'rich',
756        # ]
757    "};
758
759        assert_matches!(
760            ScriptTag::parse(contents.as_bytes()),
761            Err(Pep723Error::UnclosedBlock)
762        );
763    }
764
765    #[test]
766    fn closing_tag_trailing_whitespace() {
767        // Explicit string (not `indoc`) so the closing tag's trailing space is preserved.
768        let contents = "# /// script\n# requires-python = '>=3.11'\n# /// \n";
769
770        assert_matches!(
771            ScriptTag::parse(contents.as_bytes()),
772            Err(Pep723Error::UnclosedBlockTrailingContent)
773        );
774    }
775
776    #[test]
777    fn closing_tag_trailing_content() {
778        let contents = indoc::indoc! {r"
779            # /// script
780            # requires-python = '>=3.11'
781            # /// unexpected
782        "};
783
784        assert_matches!(
785            ScriptTag::parse(contents.as_bytes()),
786            Err(Pep723Error::UnclosedBlockTrailingContent)
787        );
788    }
789
790    #[test]
791    fn closing_tag_followed_by_prefixed_comment() {
792        let contents = indoc::indoc! {r#"
793            # /// script
794            # dependencies = []
795            # ///
796            # /// documentation
797            print("Hello, world!")
798        "#};
799
800        let actual = ScriptTag::parse(contents.as_bytes()).unwrap().unwrap();
801
802        assert_eq!(actual.metadata, "dependencies = []\n");
803        assert_eq!(
804            actual.postlude,
805            "# /// documentation\nprint(\"Hello, world!\")\n"
806        );
807    }
808
809    #[test]
810    fn closing_tag_followed_by_trailing_whitespace_comment() {
811        let contents =
812            "# /// script\n# dependencies = []\n# ///\n# /// \nprint(\"Hello, world!\")\n";
813
814        let actual = ScriptTag::parse(contents.as_bytes()).unwrap().unwrap();
815
816        assert_eq!(actual.metadata, "dependencies = []\n");
817        assert_eq!(actual.postlude, "# /// \nprint(\"Hello, world!\")\n");
818    }
819
820    #[test]
821    fn leading_content() {
822        let contents = indoc::indoc! {r"
823        pass # /// script
824        # requires-python = '>=3.11'
825        # dependencies = [
826        #   'requests<3',
827        #   'rich',
828        # ]
829        # ///
830        #
831        #
832    "};
833
834        assert_eq!(ScriptTag::parse(contents.as_bytes()).unwrap(), None);
835    }
836
837    #[test]
838    fn simple() {
839        let contents = indoc::indoc! {r"
840        # /// script
841        # requires-python = '>=3.11'
842        # dependencies = [
843        #     'requests<3',
844        #     'rich',
845        # ]
846        # ///
847
848        import requests
849        from rich.pretty import pprint
850
851        resp = requests.get('https://peps.python.org/api/peps.json')
852        data = resp.json()
853    "};
854
855        let expected_metadata = indoc::indoc! {r"
856        requires-python = '>=3.11'
857        dependencies = [
858            'requests<3',
859            'rich',
860        ]
861    "};
862
863        let expected_data = indoc::indoc! {r"
864
865        import requests
866        from rich.pretty import pprint
867
868        resp = requests.get('https://peps.python.org/api/peps.json')
869        data = resp.json()
870    "};
871
872        let actual = ScriptTag::parse(contents.as_bytes()).unwrap().unwrap();
873
874        assert_eq!(actual.prelude, String::new());
875        assert_eq!(actual.metadata, expected_metadata);
876        assert_eq!(actual.postlude, expected_data);
877    }
878
879    #[test]
880    fn simple_with_shebang() {
881        let contents = indoc::indoc! {r"
882        #!/usr/bin/env python3
883        # /// script
884        # requires-python = '>=3.11'
885        # dependencies = [
886        #     'requests<3',
887        #     'rich',
888        # ]
889        # ///
890
891        import requests
892        from rich.pretty import pprint
893
894        resp = requests.get('https://peps.python.org/api/peps.json')
895        data = resp.json()
896    "};
897
898        let expected_metadata = indoc::indoc! {r"
899        requires-python = '>=3.11'
900        dependencies = [
901            'requests<3',
902            'rich',
903        ]
904    "};
905
906        let expected_data = indoc::indoc! {r"
907
908        import requests
909        from rich.pretty import pprint
910
911        resp = requests.get('https://peps.python.org/api/peps.json')
912        data = resp.json()
913    "};
914
915        let actual = ScriptTag::parse(contents.as_bytes()).unwrap().unwrap();
916
917        assert_eq!(actual.prelude, "#!/usr/bin/env python3\n".to_string());
918        assert_eq!(actual.metadata, expected_metadata);
919        assert_eq!(actual.postlude, expected_data);
920    }
921
922    #[test]
923    fn embedded_comment() {
924        let contents = indoc::indoc! {r"
925        # /// script
926        # embedded-csharp = '''
927        # /// <summary>
928        # /// text
929        # ///
930        # /// </summary>
931        # public class MyClass { }
932        # '''
933        # ///
934    "};
935
936        let expected = indoc::indoc! {r"
937        embedded-csharp = '''
938        /// <summary>
939        /// text
940        ///
941        /// </summary>
942        public class MyClass { }
943        '''
944    "};
945
946        let actual = ScriptTag::parse(contents.as_bytes())
947            .unwrap()
948            .unwrap()
949            .metadata;
950
951        assert_eq!(actual, expected);
952    }
953
954    #[test]
955    fn trailing_lines() {
956        let contents = indoc::indoc! {r"
957            # /// script
958            # requires-python = '>=3.11'
959            # dependencies = [
960            #     'requests<3',
961            #     'rich',
962            # ]
963            # ///
964            #
965            #
966        "};
967
968        let expected = indoc::indoc! {r"
969            requires-python = '>=3.11'
970            dependencies = [
971                'requests<3',
972                'rich',
973            ]
974        "};
975
976        let actual = ScriptTag::parse(contents.as_bytes())
977            .unwrap()
978            .unwrap()
979            .metadata;
980
981        assert_eq!(actual, expected);
982    }
983
984    #[test]
985    fn unclosed_second_script_block_is_not_duplicate() {
986        let contents = indoc::indoc! {r#"
987            # /// script
988            # dependencies = ["requests"]
989            # ///
990
991            print("Hello, world!")
992
993            # /// script
994        "#};
995
996        assert!(ScriptTag::parse(contents.as_bytes()).is_ok());
997    }
998
999    #[test]
1000    fn adjacent_unclosed_second_script_block_is_not_duplicate() {
1001        let contents = indoc::indoc! {r#"
1002            # /// script
1003            # dependencies = []
1004            # ///
1005            # /// script
1006            print("Hello, world!")
1007        "#};
1008
1009        let actual = ScriptTag::parse(contents.as_bytes()).unwrap().unwrap();
1010
1011        assert_eq!(actual.metadata, "dependencies = []\n");
1012        assert_eq!(actual.postlude, "# /// script\nprint(\"Hello, world!\")\n");
1013    }
1014
1015    #[test]
1016    fn other_script_block_is_ignored() {
1017        let contents = indoc::indoc! {r#"
1018            # /// script
1019            # dependencies = ["requests"]
1020            # ///
1021
1022
1023            # /// other
1024            # /// script
1025            # ///
1026
1027            print("Hello, world!")
1028        "#};
1029
1030        assert!(ScriptTag::parse(contents.as_bytes()).is_ok());
1031    }
1032
1033    #[test]
1034    fn serialize_metadata_formatting() {
1035        let metadata = indoc::indoc! {r"
1036            requires-python = '>=3.11'
1037            dependencies = [
1038              'requests<3',
1039              'rich',
1040            ]
1041        "};
1042
1043        let expected_output = indoc::indoc! {r"
1044            # /// script
1045            # requires-python = '>=3.11'
1046            # dependencies = [
1047            #   'requests<3',
1048            #   'rich',
1049            # ]
1050            # ///
1051        "};
1052
1053        let result = serialize_metadata(metadata);
1054        assert_eq!(result, expected_output);
1055    }
1056
1057    #[test]
1058    fn serialize_metadata_empty() {
1059        let metadata = "";
1060        let expected_output = "# /// script\n# ///\n";
1061
1062        let result = serialize_metadata(metadata);
1063        assert_eq!(result, expected_output);
1064    }
1065
1066    #[test]
1067    fn script_init_empty() {
1068        let contents = "".as_bytes();
1069        let (prelude, metadata, postlude) =
1070            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1071                .unwrap();
1072        assert_eq!(prelude, "");
1073        assert_eq!(
1074            metadata.raw,
1075            indoc::indoc! {r"
1076            dependencies = []
1077            "}
1078        );
1079        assert_eq!(postlude, "");
1080    }
1081
1082    #[test]
1083    fn script_init_requires_python() {
1084        let contents = "".as_bytes();
1085        let (prelude, metadata, postlude) = Pep723Script::init_metadata(
1086            contents,
1087            &uv_pep440::VersionSpecifiers::from_str(">=3.8").unwrap(),
1088        )
1089        .unwrap();
1090        assert_eq!(prelude, "");
1091        assert_eq!(
1092            metadata.raw,
1093            indoc::indoc! {r#"
1094            requires-python = ">=3.8"
1095            dependencies = []
1096            "#}
1097        );
1098        assert_eq!(postlude, "");
1099    }
1100
1101    #[test]
1102    fn script_init_with_hashbang() {
1103        let contents = indoc::indoc! {r#"
1104        #!/usr/bin/env python3
1105
1106        print("Hello, world!")
1107        "#}
1108        .as_bytes();
1109        let (prelude, metadata, postlude) =
1110            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1111                .unwrap();
1112        assert_eq!(prelude, "#!/usr/bin/env python3\n");
1113        assert_eq!(
1114            metadata.raw,
1115            indoc::indoc! {r"
1116            dependencies = []
1117            "}
1118        );
1119        assert_eq!(
1120            postlude,
1121            indoc::indoc! {r#"
1122
1123            print("Hello, world!")
1124            "#}
1125        );
1126    }
1127
1128    #[test]
1129    fn script_init_with_other_metadata() {
1130        let contents = indoc::indoc! {r#"
1131        # /// noscript
1132        # Hello,
1133        #
1134        # World!
1135        # ///
1136
1137        print("Hello, world!")
1138        "#}
1139        .as_bytes();
1140        let (prelude, metadata, postlude) =
1141            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1142                .unwrap();
1143        assert_eq!(prelude, "");
1144        assert_eq!(
1145            metadata.raw,
1146            indoc::indoc! {r"
1147            dependencies = []
1148            "}
1149        );
1150        // Note the extra line at the beginning.
1151        assert_eq!(
1152            postlude,
1153            indoc::indoc! {r#"
1154
1155            # /// noscript
1156            # Hello,
1157            #
1158            # World!
1159            # ///
1160
1161            print("Hello, world!")
1162            "#}
1163        );
1164    }
1165
1166    #[test]
1167    fn script_init_with_hashbang_and_other_metadata() {
1168        let contents = indoc::indoc! {r#"
1169        #!/usr/bin/env python3
1170        # /// noscript
1171        # Hello,
1172        #
1173        # World!
1174        # ///
1175
1176        print("Hello, world!")
1177        "#}
1178        .as_bytes();
1179        let (prelude, metadata, postlude) =
1180            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1181                .unwrap();
1182        assert_eq!(prelude, "#!/usr/bin/env python3\n");
1183        assert_eq!(
1184            metadata.raw,
1185            indoc::indoc! {r"
1186            dependencies = []
1187            "}
1188        );
1189        // Note the extra line at the beginning.
1190        assert_eq!(
1191            postlude,
1192            indoc::indoc! {r#"
1193
1194            # /// noscript
1195            # Hello,
1196            #
1197            # World!
1198            # ///
1199
1200            print("Hello, world!")
1201            "#}
1202        );
1203    }
1204
1205    #[test]
1206    fn script_init_with_valid_metadata_line() {
1207        let contents = indoc::indoc! {r#"
1208        # Hello,
1209        # /// noscript
1210        #
1211        # World!
1212        # ///
1213
1214        print("Hello, world!")
1215        "#}
1216        .as_bytes();
1217        let (prelude, metadata, postlude) =
1218            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1219                .unwrap();
1220        assert_eq!(prelude, "");
1221        assert_eq!(
1222            metadata.raw,
1223            indoc::indoc! {r"
1224            dependencies = []
1225            "}
1226        );
1227        // Note the extra line at the beginning.
1228        assert_eq!(
1229            postlude,
1230            indoc::indoc! {r#"
1231
1232            # Hello,
1233            # /// noscript
1234            #
1235            # World!
1236            # ///
1237
1238            print("Hello, world!")
1239            "#}
1240        );
1241    }
1242
1243    #[test]
1244    fn script_init_with_valid_empty_metadata_line() {
1245        let contents = indoc::indoc! {r#"
1246        #
1247        # /// noscript
1248        # Hello,
1249        # World!
1250        # ///
1251
1252        print("Hello, world!")
1253        "#}
1254        .as_bytes();
1255        let (prelude, metadata, postlude) =
1256            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1257                .unwrap();
1258        assert_eq!(prelude, "");
1259        assert_eq!(
1260            metadata.raw,
1261            indoc::indoc! {r"
1262            dependencies = []
1263            "}
1264        );
1265        // Note the extra line at the beginning.
1266        assert_eq!(
1267            postlude,
1268            indoc::indoc! {r#"
1269
1270            #
1271            # /// noscript
1272            # Hello,
1273            # World!
1274            # ///
1275
1276            print("Hello, world!")
1277            "#}
1278        );
1279    }
1280
1281    #[test]
1282    fn script_init_with_non_metadata_comment() {
1283        let contents = indoc::indoc! {r#"
1284        #Hello,
1285        # /// noscript
1286        #
1287        # World!
1288        # ///
1289
1290        print("Hello, world!")
1291        "#}
1292        .as_bytes();
1293        let (prelude, metadata, postlude) =
1294            Pep723Script::init_metadata(contents, &uv_pep440::VersionSpecifiers::default())
1295                .unwrap();
1296        assert_eq!(prelude, "");
1297        assert_eq!(
1298            metadata.raw,
1299            indoc::indoc! {r"
1300            dependencies = []
1301            "}
1302        );
1303        assert_eq!(
1304            postlude,
1305            indoc::indoc! {r#"
1306            #Hello,
1307            # /// noscript
1308            #
1309            # World!
1310            # ///
1311
1312            print("Hello, world!")
1313            "#}
1314        );
1315    }
1316}