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