Skip to main content

uv_scripts/
lib.rs

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