dot_writer/

writer.rs

use std::io::Write;

use super::scope::Scope;

const INDENT_STEP: usize = 2;

/// The entry point struct for writing a DOT graph.
/// See the examples on the index for how to use it.
/// This struct must live for the livetime of writing the graph,
/// and also outlive the [`Write`] struct it points to,
/// because it's [`Drop`] trait will finish writing graph.
pub struct DotWriter<'w> {
    writer: &'w mut dyn Write,
    indent: usize,
    next_id: usize,
    digraph: bool,
    pretty_print: bool,
}

/// This is the only constructor, which requires a [`Write`]
/// to borrow and write the DOT output to.
/// Defaults to pretty printing the output (see [`DotWriter::set_pretty_print`]).
impl<'w, W: Write> From<&'w mut W> for DotWriter<'w> {
    fn from(writer: &'w mut W) -> Self {
        Self {
            writer,
            indent: 0,
            next_id: 0,
            digraph: false,
            pretty_print: true,
        }
    }
}

impl<'w> DotWriter<'w> {
    /// If set to true, then output will have additional whitespace
    /// including newlines and indentation to make the output more
    /// human friendly. If false then this will be left out and the output
    /// will be more compact. Defaults to true. For example disabled:
    ///
    /// ```
    /// use dot_writer::DotWriter;
    ///
    /// let mut bytes = Vec::new();
    /// let mut writer = DotWriter::from(&mut bytes);
    /// writer.set_pretty_print(false);
    /// writer.graph().subgraph().cluster().subgraph();
    /// println!("{}", std::str::from_utf8(&bytes).unwrap());
    /// ```
    ///
    /// Gives you:
    ///
    /// `
    /// graph {subgraph {subgraph cluster_0 {subgraph {}}}}
    /// `
    ///
    /// While enabled:
    ///
    /// ```
    /// use dot_writer::DotWriter;
    ///
    /// let mut bytes = Vec::new();
    /// let mut writer = DotWriter::from(&mut bytes);
    /// // writer.set_pretty_print(false); defaults to true anyway
    /// writer.graph().subgraph().cluster().subgraph();
    /// println!("{}", std::str::from_utf8(&bytes).unwrap());
    /// ```
    ///
    /// Gives you:
    ///
    /// ```txt
    /// graph {
    ///   subgraph {
    ///     subgraph cluster_0 {
    ///       subgraph {}
    ///     }
    ///   }
    /// }
    ///```
    pub fn set_pretty_print(&mut self, pretty_print: bool) {
        self.pretty_print = pretty_print;
    }

    /// Start a new undirection graph. This uses the edge operator `--`
    /// automatically, which means edges should be rendered without any
    /// arrowheads by default.    
    /// ```
    /// use dot_writer::DotWriter;
    ///
    /// let mut bytes = Vec::new();
    /// let mut writer = DotWriter::from(&mut bytes);
    /// writer.set_pretty_print(false);
    /// writer.graph().edge("a", "b");
    /// assert_eq!(
    ///     std::str::from_utf8(&bytes).unwrap(),
    ///     "graph{a--b;}"
    /// );
    /// ```
    pub fn graph(&mut self) -> Scope<'_, 'w> {
        self.digraph = false;
        Scope::new(self, b"graph")
    }

    /// Start a new directed graph. This uses the edge operator `->`
    /// automatically, which means edges should be redered with an
    /// arrow head pointing to the head (end) node.    
    /// ```
    /// use dot_writer::DotWriter;
    ///
    /// let mut bytes = Vec::new();
    /// let mut writer = DotWriter::from(&mut bytes);
    /// writer.set_pretty_print(false);
    /// writer.digraph().edge("a", "b");
    /// assert_eq!(
    ///     std::str::from_utf8(&bytes).unwrap(),
    ///     "digraph{a->b;}"
    /// );
    /// ```
    pub fn digraph(&mut self) -> Scope<'_, 'w> {
        self.digraph = true;
        Scope::new(self, b"digraph")
    }

    /// Uses a callback to write DOT to a [`String`]. This is useful
    /// if you just want to write your dot code to a string rather than
    /// a file or stdout, and want less boiler plate setup code.
    /// It's used internally for unit testing, so the output is not
    /// pretty printed by default (but you can overwrite that by calling
    /// [`DotWriter::set_pretty_print`] from within the callback).
    ///
    /// ```
    /// use dot_writer::DotWriter;
    ///
    /// let output = DotWriter::write_string(|writer| {
    ///     let mut graph = writer.graph();
    ///     graph.edge("a", "b");
    /// });
    /// assert_eq!(output, "graph{a--b;}");
    /// ```
    pub fn write_string<F: Fn(&mut DotWriter)>(builder: F) -> String {
        let mut bytes = Vec::new();
        let mut writer = DotWriter::from(&mut bytes);
        writer.set_pretty_print(false);
        (builder)(&mut writer);
        String::from_utf8(bytes).unwrap()
    }
}

impl<'w> DotWriter<'w> {
    pub(crate) fn write(&mut self, bytes: &[u8]) {
        if self.pretty_print {
            self.writer.write_all(bytes).unwrap();
        } else {
            for b in bytes.iter().filter(|b| !b.is_ascii_whitespace()) {
                self.writer.write_all(&[*b]).unwrap();
            }
        }
    }

    pub(crate) fn write_with_whitespace(&mut self, bytes: &[u8]) {
        self.writer.write_all(bytes).unwrap();
    }

    pub(crate) fn write_quoted(&mut self, bytes: &[u8]) {
        self.writer.write_all(b"\"").unwrap();
        self.write_with_whitespace(bytes);
        self.writer.write_all(b"\"").unwrap();
    }

    pub(crate) fn write_indent(&mut self) {
        for _ in 0..self.indent {
            self.writer.write_all(b" ").unwrap();
        }
    }

    pub(crate) fn write_edge_operator(&mut self) {
        match self.digraph {
            true => self.write(b" -> "),
            false => self.write(b" -- "),
        }
    }

    pub(crate) fn indent(&mut self) {
        if self.pretty_print {
            self.indent += INDENT_STEP;
        }
    }

    pub(crate) fn unindent(&mut self) {
        if self.pretty_print {
            self.indent -= INDENT_STEP;
        }
    }

    pub(crate) fn next_id(&mut self) -> usize {
        let next_id = self.next_id;
        self.next_id += 1;
        next_id
    }
}

//
// Statement
//

pub(crate) struct Statement<'d, 'w> {
    line: Line<'d, 'w>,
}

impl<'d, 'w> Statement<'d, 'w> {
    pub(crate) fn new(writer: &'d mut DotWriter<'w>) -> Self {
        Self {
            line: Line::new(writer),
        }
    }
}

impl<'d, 'w> std::ops::DerefMut for Statement<'d, 'w> {
    fn deref_mut(&mut self) -> &mut DotWriter<'w> {
        &mut *self.line
    }
}

impl<'d, 'w> std::ops::Deref for Statement<'d, 'w> {
    type Target = DotWriter<'w>;
    fn deref(&self) -> &DotWriter<'w> {
        &*self.line
    }
}

impl<'d, 'w> Drop for Statement<'d, 'w> {
    fn drop(&mut self) {
        self.line.write(b";");
    }
}

//
// Line
//

pub(crate) struct Line<'d, 'w> {
    writer: &'d mut DotWriter<'w>,
}

impl<'d, 'w> Line<'d, 'w> {
    pub(crate) fn new(writer: &'d mut DotWriter<'w>) -> Self {
        writer.write_indent();
        Self { writer }
    }
}

impl<'d, 'w> std::ops::DerefMut for Line<'d, 'w> {
    fn deref_mut(&mut self) -> &mut DotWriter<'w> {
        self.writer
    }
}

impl<'d, 'w> std::ops::Deref for Line<'d, 'w> {
    type Target = DotWriter<'w>;
    fn deref(&self) -> &DotWriter<'w> {
        self.writer
    }
}

impl<'a, 'b: 'a> Drop for Line<'a, 'b> {
    fn drop(&mut self) {
        if self.writer.pretty_print {
            self.writer.write(b"\n");
        }
    }
}