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