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