Skip to main content

uv_scripts/
lib.rs

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