plantuml_server_client_rs/

locate.rs

use crate::Format;
use anyhow::{Context as _, Result};
use bytes::Bytes;
use futures::{Stream, StreamExt};
use std::path::{Path, PathBuf};
use tokio::io::{AsyncWrite, BufWriter};

/// An wrapper for representing file or stdin / stdout
#[derive(Debug)]
pub enum Locate {
    File(PathBuf),
    StdInOut,
}

impl Locate {
    /// Creates a new [`Locate`].
    /// - Some(x) -> [`Locate::File`]
    /// - None -> [`Locate::StdInOut`] (stdin / stdout)
    pub fn new(locate: Option<PathBuf>) -> Self {
        match locate {
            Some(locate) if locate.to_string_lossy() == "-" => Self::StdInOut,
            Some(locate) => Self::File(locate),
            None => Self::StdInOut,
        }
    }

    /// Converts a [`Locate`] to an owned [`PathBuf`].
    pub fn to_path_buf(&self) -> PathBuf {
        match self {
            Self::File(input) => input.clone(),
            Self::StdInOut => PathBuf::from("."),
        }
    }

    /// Creates image's output path.
    ///
    /// * `input` - An input [`Locate`].
    /// * `id` - Identifiers in the PlantUML diagram.
    /// * `format` - A format of output.
    ///
    /// # Examples
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate, Format};
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() {
    /// let locate = Locate::from("./output/dir");
    /// let output_path = locate.output_path(&"./foo.puml".into(), "id_a", &Format::Svg);
    ///
    /// let expected = Locate::from(PathBuf::from("./output/dir/foo/id_a.svg"));
    /// assert_eq!(output_path, expected);
    /// # }
    /// ```
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate, Format};
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() {
    /// let locate = Locate::from(None); // `--output` is not specified
    /// let output_path = locate.output_path(&"./foo.puml".into(), "id_a", &Format::Svg);
    ///
    /// let expected = Locate::from(PathBuf::from("foo/id_a.svg"));
    /// assert_eq!(output_path, expected);
    /// # }
    /// ```
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate, Format};
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() {
    /// let locate = Locate::from("./output/dir");
    /// let input = Locate::from(None); // `--input` is not specified (read from stdin)
    /// let output_path = locate.output_path(&input, "id_a", &Format::Svg);
    ///
    /// let expected = Locate::from(PathBuf::from("./output/dir/id_a.svg"));
    /// assert_eq!(output_path, expected);
    /// # }
    /// ```
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate, Format};
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() {
    /// let locate = Locate::from(None);
    /// let input = Locate::from(None);
    /// let output_path = locate.output_path(&input, "N/A", &Format::Svg);
    ///
    /// let expected = Locate::from(None);
    /// assert_eq!(output_path, expected);
    /// # }
    /// ```
    pub fn output_path(&self, input: &Locate, id: &str, format: &Format) -> Self {
        let extension: &str = format.into();
        self.output_path_inner(input, id, extension)
    }

    /// Creates "combined" PlantUML content's path.
    /// The "combined PlantUML content" is output when the `--combined` option is specified.
    ///
    /// # Examples
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate, Format};
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() {
    /// let locate = Locate::from("./output/dir");
    /// let output_path = locate.combined_output_path(&"./foo.puml".into(), "id_a");
    ///
    /// let expected_path = "./output/dir/foo/id_a.puml"; // different extension from `output_path()`.
    /// let expected = Some(Locate::from(PathBuf::from(expected_path)));
    /// assert_eq!(output_path, expected);
    /// # }
    /// ```
    pub fn combined_output_path(&self, input: &Locate, id: &str) -> Option<Self> {
        let path = self.output_path_inner(input, id, "puml");

        match path {
            Self::StdInOut => None,
            Self::File(file) => Some(Self::File(file)),
        }
    }

    fn output_path_inner(&self, input: &Locate, id: &str, extension: &str) -> Self {
        if self.is_std_inout() && input.is_std_inout() {
            return Self::StdInOut;
        }

        let input = match input {
            Self::File(x) => {
                let x = x.clone();
                match x.clone().file_name().map(PathBuf::from) {
                    Some(mut x) => {
                        x.set_extension("");
                        x
                    }
                    None => PathBuf::new(),
                }
            }
            Self::StdInOut => PathBuf::new(),
        };

        let mut base = match self {
            Self::File(x) => x.clone(),
            Self::StdInOut => PathBuf::new(),
        };

        base.push(input);
        base.push(id);
        base.set_extension(extension);

        Self::File(base)
    }

    /// Read from the [`Locate`].
    ///
    /// # Examples
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate };
    /// # use std::path::PathBuf;
    /// #
    /// # fn main() -> anyhow::Result<()> {
    /// # let _ = std::fs::create_dir_all("./target/tests_docs/");
    /// let locate = Locate::from("./target/tests_docs/locate_read");
    /// let data = "foo";
    ///
    /// std::fs::write(locate.to_path_buf(), data)?;
    /// let read_data = locate.read()?;
    /// assert_eq!(read_data, data);
    /// # Ok(())
    /// # }
    /// ```
    pub fn read(&self) -> Result<String> {
        use std::fs::read_to_string;

        match self {
            Self::File(input) => read_to_string(input)
                .with_context(|| format!("failed to read file: input = ({input:?})")),

            Self::StdInOut => Self::read_from_stdin().context("failed to read stdin"),
        }
    }

    fn read_from_stdin() -> Result<String> {
        let mut delimiter = "";

        std::io::stdin()
            .lines()
            .try_fold(String::new(), |mut acc, item| {
                acc += delimiter;
                acc += &item?;
                delimiter = "\n";
                Ok(acc)
            })
    }

    /// Write a stream to [`Locate`].
    ///
    /// * `stream` - Data to write
    ///
    /// # Examples
    ///
    /// ```
    /// # use plantuml_server_client_rs as pscr;
    /// #
    /// # use pscr::{Locate };
    /// # use std::path::PathBuf;
    /// #
    /// # #[tokio::main]
    /// # async fn main() -> anyhow::Result<()> {
    /// # let _ = std::fs::create_dir_all("./target/tests_docs/");
    /// let locate = Locate::from("./target/tests_docs/locate_write");
    /// let data = "foo2";
    /// let stream = futures::stream::iter(vec![Ok(data.into())]);
    ///
    /// locate.write(stream).await?;
    /// let read_data = std::fs::read_to_string(locate.to_path_buf())?;
    /// assert_eq!(read_data, data);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn write(
        &self,
        stream: impl Stream<Item = Result<Bytes, reqwest::Error>> + Unpin,
    ) -> Result<()> {
        use tokio::fs::File;

        match self {
            Self::File(output) => {
                tracing::info!("output = {:?}", output);

                create_base_dir(&output).await?;

                let file = File::create(output)
                    .await
                    .with_context(|| format!("failed to open write file: output = {output:?}"))?;

                let writer = BufWriter::new(file);
                write_inner(writer, stream).await
            }

            Self::StdInOut => {
                let writer = BufWriter::new(tokio::io::stdout());
                write_inner(writer, stream).await
            }
        }
    }

    fn is_std_inout(&self) -> bool {
        matches!(self, Self::StdInOut)
    }
}

impl From<Option<PathBuf>> for Locate {
    fn from(inner: Option<PathBuf>) -> Self {
        Self::new(inner)
    }
}

impl From<PathBuf> for Locate {
    fn from(inner: PathBuf) -> Self {
        Self::new(Some(inner))
    }
}

impl From<&str> for Locate {
    fn from(inner: &str) -> Self {
        Self::new(Some(PathBuf::from(inner)))
    }
}

impl PartialEq<Locate> for Locate {
    fn eq(&self, rhs: &Locate) -> bool {
        match (self, rhs) {
            (Self::File(lhs), Self::File(rhs)) => lhs.eq(rhs),
            (Self::StdInOut, Self::StdInOut) => true,
            _ => false,
        }
    }
}

async fn write_inner<W>(
    mut writer: BufWriter<W>,
    mut stream: impl Stream<Item = Result<Bytes, reqwest::Error>> + Unpin,
) -> Result<()>
where
    W: AsyncWrite + Unpin,
{
    use tokio::io::AsyncWriteExt;
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.context("failed to read chunk")?;
        writer.write_all(&chunk).await?;
    }
    writer.flush().await?;
    Ok(())
}

async fn create_base_dir<P: AsRef<Path>>(output: P) -> Result<()> {
    use tokio::fs::create_dir_all;

    let mut base = output.as_ref().to_path_buf();
    base.pop();
    create_dir_all(base)
        .await
        .with_context(|| "failed to create directory: directory = {base:?}")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_output_path() {
        let format = Format::Svg;

        // in: Some(_), out: Some(_)
        let input = Locate::new(Some(PathBuf::from("input_dir/input_1.puml")));
        let output = Locate::new(Some(PathBuf::from("output_dir")));

        let path = output.output_path(&input, "id_1", &format);
        assert_eq!(
            path.to_path_buf(),
            PathBuf::from("output_dir/input_1/id_1.svg")
        );
        let path = output.output_path(&input, "id_2", &format);
        assert_eq!(
            path.to_path_buf(),
            PathBuf::from("output_dir/input_1/id_2.svg")
        );

        // in: Some(_), out: None
        let input = Locate::new(Some(PathBuf::from("input_dir/input_2.puml")));
        let output = Locate::new(None);

        let path = output.output_path(&input, "id_1", &format);
        assert_eq!(path.to_path_buf(), PathBuf::from("input_2/id_1.svg"));

        // in: None, out: Some(_)
        let input = Locate::new(None);
        let output = Locate::new(Some(PathBuf::from("output_dir")));

        let path = output.output_path(&input, "id_1", &format);
        assert_eq!(path.to_path_buf(), PathBuf::from("output_dir/id_1.svg"));

        // in: None, out: None
        let input = Locate::new(None);
        let output = Locate::new(None);

        let path = output.output_path(&input, "id_1", &format);
        assert_eq!(path.to_path_buf(), PathBuf::from("."));
    }
}