Skip to main content

proto_build_kit/
stage.rs

1// SPDX-License-Identifier: MIT
2//! Stage embedded proto bytes onto a tempdir at protoc-relative paths.
3
4use std::collections::BTreeSet;
5use std::io::Write;
6use std::path::PathBuf;
7
8use crate::Error;
9
10/// Accumulates `(relative_path, bytes)` pairs and writes them to a
11/// fresh tempdir laid out as `<tmp>/<relative_path>`.
12///
13/// `relative_path` is the protoc-style import path the consuming
14/// `.proto` will use (`bones/v1/pagination.proto`, etc.). The
15/// returned [`tempfile::TempDir`] cleans up on drop — hold it until
16/// codegen finishes.
17///
18/// # Examples
19///
20/// ```no_run
21/// use proto_build_kit::Stager;
22///
23/// const FOO: &[u8] = b"syntax = \"proto3\"; package foo.v1; message Foo {}";
24///
25/// fn build() -> Result<(), Box<dyn std::error::Error>> {
26///     let staged = Stager::new()
27///         .add("foo/v1/foo.proto", FOO)
28///         .stage()?;
29///     // Use staged.path() on your protoc include path.
30///     Ok(())
31/// }
32/// ```
33///
34/// Pair with a sibling `*-protos` crate that exposes `files()`:
35///
36/// ```ignore
37/// let staged = Stager::new()
38///     .with(some_protos::files())
39///     .with(other_protos::files())
40///     .stage()?;
41/// ```
42///
43/// # Duplicate detection
44///
45/// Two entries with the same `relative_path` cause `stage()` to return
46/// [`Error::DuplicatePath`]. This protects against silent shadowing when
47/// two `*-protos` crates collide on a path — almost always a bug.
48#[derive(Default)]
49pub struct Stager {
50    files: Vec<(&'static str, &'static [u8])>,
51}
52
53impl Stager {
54    /// Construct an empty stager.
55    #[must_use]
56    pub fn new() -> Self {
57        Self::default()
58    }
59
60    /// Append a single `(relative_path, bytes)` entry.
61    #[must_use]
62    pub fn add(mut self, relative_path: &'static str, bytes: &'static [u8]) -> Self {
63        self.files.push((relative_path, bytes));
64        self
65    }
66
67    /// Append every entry yielded by `iter`. Typically called with a
68    /// `*-protos` crate's `files()` accessor.
69    #[must_use]
70    pub fn with<I>(mut self, iter: I) -> Self
71    where
72        I: IntoIterator<Item = (&'static str, &'static [u8])>,
73    {
74        self.files.extend(iter);
75        self
76    }
77
78    /// Write every staged entry to a fresh tempdir and return the
79    /// handle. Add `tempdir.path()` to your protoc include path.
80    ///
81    /// # Errors
82    ///
83    /// - [`Error::DuplicatePath`] if two entries declare the same
84    ///   relative path.
85    /// - [`Error::Io`] if tempdir creation, directory creation, or
86    ///   file writes fail.
87    pub fn stage(self) -> Result<tempfile::TempDir, Error> {
88        let mut seen: BTreeSet<&str> = BTreeSet::new();
89        for (path, _) in &self.files {
90            if !seen.insert(*path) {
91                return Err(Error::DuplicatePath {
92                    path: (*path).to_string(),
93                });
94            }
95        }
96
97        let dir = tempfile::tempdir()?;
98        for (rel, bytes) in self.files {
99            let target: PathBuf = dir.path().join(rel);
100            if let Some(parent) = target.parent() {
101                std::fs::create_dir_all(parent)?;
102            }
103            let mut f = std::fs::File::create(&target)?;
104            f.write_all(bytes)?;
105        }
106        Ok(dir)
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn stages_single_file_at_relative_path() {
116        let dir = Stager::new()
117            .add("a/v1/x.proto", b"syntax = \"proto3\";")
118            .stage()
119            .expect("stage");
120        let p = dir.path().join("a/v1/x.proto");
121        assert!(p.exists(), "expected staged file at {p:?}");
122        let body = std::fs::read(&p).expect("read");
123        assert_eq!(body.as_slice(), b"syntax = \"proto3\";");
124    }
125
126    #[test]
127    fn stages_multiple_files_under_nested_subdirs() {
128        let dir = Stager::new()
129            .add("a/v1/x.proto", b"x")
130            .add("b/v2/y.proto", b"y")
131            .stage()
132            .expect("stage");
133        assert!(dir.path().join("a/v1/x.proto").exists());
134        assert!(dir.path().join("b/v2/y.proto").exists());
135    }
136
137    #[test]
138    fn with_iterator_extends_files() {
139        let pairs: Vec<(&str, &[u8])> = vec![
140            ("a/v1/x.proto", b"x" as &[u8]),
141            ("b/v1/y.proto", b"y" as &[u8]),
142        ];
143        let dir = Stager::new().with(pairs).stage().expect("stage");
144        assert!(dir.path().join("a/v1/x.proto").exists());
145        assert!(dir.path().join("b/v1/y.proto").exists());
146    }
147
148    #[test]
149    fn duplicate_paths_return_error() {
150        let err = Stager::new()
151            .add("a/v1/x.proto", b"first")
152            .add("a/v1/x.proto", b"second")
153            .stage()
154            .expect_err("should error on duplicate");
155        match err {
156            Error::DuplicatePath { path } => assert_eq!(path, "a/v1/x.proto"),
157            other => panic!("expected DuplicatePath, got {other:?}"),
158        }
159    }
160}