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