add custom attribute adding function

src/edge.rs

@@ -1,52 +1,41 @@
-use crate::{
-    arrow::{Arrow},
-    style::{Style},
-    utils::{quote_string},
-};
+use crate::{arrow::Arrow, style::Style, utils::quote_string};
 
 /// `Graph`'s edge.
 #[derive(Clone)]
 pub struct Edge {
     from: String,
     to: String,
-    label: String,
-    label_url: String,
-    url: String,
-    style: Style,
     start_arrow: Arrow,
     end_arrow: Arrow,
-    color: Option<String>,
+    attribs: Vec<String>,
 }
 
 impl Edge {
-    pub fn new(from: &str, to: &str, label: &str) -> Self {
-        Edge { 
-            from: String::from(from), to: String::from(to), 
-            label: String::from(label), label_url: Default::default(),
-            color: None, style: Style::None, 
-            start_arrow: Arrow::default(), end_arrow: Arrow::default(),
-            url: Default::default() 
+    pub fn new(from: &str, to: &str) -> Self {
+        Edge {
+            from: String::from(from),
+            to: String::from(to),
+            start_arrow: Arrow::default(),
+            end_arrow: Arrow::default(),
+            attribs: vec![],
         }
     }
 
     pub fn label(&mut self, label: &str) -> Self {
-        let mut edge = self.clone();
-        edge.label = String::from(label);
-        edge
+        self.attrib("label", label)
     }
 
     pub fn style(&mut self, style: Style) -> Self {
-        let mut edge = self.clone();
-        edge.style = style;
-        edge
+        self.attrib("style", style.as_slice())
     }
 
-    pub fn color(&mut self, color: Option<&str>) -> Self {
+    pub fn color(&mut self, color: &str) -> Self {
+        self.attrib("color", &quote_string(color.to_owned()))
+    }
+
+    pub fn attrib(&self, name: &str, value: &str) -> Self {
         let mut edge = self.clone();
-        edge.color = match color {
-            Some(c) => Some(String::from(c)),
-            None => None
-        };
+        edge.attribs.push(format!("{}={}", name, value));
         edge
     }
 
@@ -62,73 +51,52 @@ impl Edge {
         edge
     }
 
-    pub fn label_url(&mut self, url: String) -> Self {
-        let mut edge = self.clone();
-        edge.label_url = url;
-        edge
+    pub fn label_url(&mut self, url: &str) -> Self {
+        self.attrib("labelURL", &quote_string(url.to_owned()))
     }
 
-    pub fn url(&mut self, url: String) -> Self {
-        let mut edge = self.clone();
-        edge.url = url;
-        edge
+    pub fn url(&mut self, url: &str) -> Self {
+        self.attrib("URL", &quote_string(url.to_owned()))
     }
 
     pub fn to_dot_string(&self, edge_symbol: &str) -> String {
-        let colorstring: String;
-        let escaped_label: &String = &quote_string(self.label.clone());
         let start_arrow_s: String = self.start_arrow.to_dot_string();
         let end_arrow_s: String = self.end_arrow.to_dot_string();
-        let escaped_label_url: &String = &quote_string(self.label_url.clone());
-        let escaped_url: &String = &quote_string(self.url.clone());
-
-        let mut text = vec!["\"", self.from.as_str(), "\" ",
-        edge_symbol, " ",
-        "\"", self.to.as_str(), "\"",];
-
-        text.push("[label=");
-        text.push(escaped_label.as_str());
-        text.push("]");
-
-    if !self.label_url.is_empty(){
-        text.push("[labelURL=");
-        text.push(escaped_label_url.as_str());
-        text.push("]");
-    }
-
-    if !self.url.is_empty(){
-        text.push("[URL=");
-        text.push(escaped_url.as_str());
-        text.push("]");
-    }
 
-        if self.style != Style::None {
-            text.push("[style=\"");
-            text.push(self.style.as_slice());
-            text.push("\"]");
-        }
+        let mut text = vec![
+            "\"",
+            self.from.as_str(),
+            "\" ",
+            edge_symbol,
+            " ",
+            "\"",
+            self.to.as_str(),
+            "\"",
+        ];
 
-        let color: Option<String> = match &self.color {
-            Some(l) => {
-                Some((*l).clone())
-            },
-            None => None,
-        };
-        if let Some(c) = color {
-            colorstring = quote_string(c);
-            text.push("[color=");
-            text.push(&colorstring);
+        let binding = self.attribs.join(",");
+        if !self.attribs.is_empty() {
+            text.push("[");
+            text.push(binding.as_str());
             text.push("]");
         }
 
         let mut arrow_text: Vec<String> = vec![];
         let mut arrow_str: String = String::new();
         if !self.start_arrow.is_default() || !self.end_arrow.is_default() {
             if !self.end_arrow.is_default() {
-                arrow_text.push(vec!["arrowhead=\"", &end_arrow_s, "\""].into_iter().collect());
+                arrow_text.push(
+                    vec!["arrowhead=\"", &end_arrow_s, "\""]
+                        .into_iter()
+                        .collect(),
+                );
             }
             if !self.start_arrow.is_default() {
-                arrow_text.push(vec!["arrowtail=\"", &start_arrow_s, "\""].into_iter().collect());
+                arrow_text.push(
+                    vec!["arrowtail=\"", &start_arrow_s, "\""]
+                        .into_iter()
+                        .collect(),
+                );
             }
             if !self.start_arrow.is_default() && !self.end_arrow.is_default() {
                 arrow_text.push(String::from("dir=\"both\""));
@@ -140,8 +108,7 @@ impl Edge {
             arrow_str.push_str("]");
             text.push(arrow_str.as_str());
         }
-
         text.push(";");
         return text.into_iter().collect();
     }
-}
+}

src/graph.rs

@@ -1,9 +1,6 @@
-use crate::{
-    node::{Node},
-    edge::{Edge}, subgraph::Subgraph, utils::quote_string,
-};
-use std::io::prelude::*;
+use crate::{edge::Edge, node::Node, subgraph::Subgraph, utils::quote_string};
 use std::io;
+use std::io::prelude::*;
 
 /// Entry point of this library, use `to_dot_string` to get the string output.
 #[derive(Clone)]
@@ -13,12 +10,21 @@ pub struct Graph {
     url: String,
     nodes: Vec<Node>,
     edges: Vec<Edge>,
-    subgraph: Vec<Subgraph>
+    subgraph: Vec<Subgraph>,
+    attribs: Vec<String>,
 }
 
 impl Graph {
     pub fn new(name: &str, kind: Kind) -> Graph {
-        Graph { name: String::from(name), kind: kind, nodes: vec![], edges: vec![], subgraph: vec![], url: Default::default() }
+        Graph {
+            name: String::from(name),
+            kind: kind,
+            nodes: vec![],
+            edges: vec![],
+            subgraph: vec![],
+            url: Default::default(),
+            attribs: vec![],
+        }
     }
 
     pub fn add_node(&mut self, node: Node) -> () {
@@ -47,9 +53,15 @@ impl Graph {
         Ok(s)
     }
 
+    pub fn attrib(&self, name: &str, value: &str) -> Self {
+        let mut graph = self.clone();
+        graph.attribs.push(format!("{}={}", name, value));
+        graph
+    }
+
     /// Renders graph `g` into the writer `w` in DOT syntax.
     /// (Main entry point for the library.)
-    fn render_opts<'a,W: Write>(&self, w: &mut W) -> io::Result<()> {
+    fn render_opts<'a, W: Write>(&self, w: &mut W) -> io::Result<()> {
         fn writeln<W: Write>(w: &mut W, arg: &[&str]) -> io::Result<()> {
             for &s in arg {
                 w.write_all(s.as_bytes())?;
@@ -63,7 +75,12 @@ impl Graph {
 
         writeln(w, &[self.kind.keyword(), " ", self.name.as_str(), " {"])?;
 
-        if !self.url.is_empty(){
+        for attrib in &self.attribs {
+            indent(w)?;
+            writeln(w, &[&attrib, ";"])?;
+        }
+
+        if !self.url.is_empty() {
             indent(w)?;
             writeln(w, &["URL=", quote_string(self.url.clone()).as_str()])?;
         }
@@ -111,7 +128,7 @@ impl Kind {
     pub fn keyword(&self) -> &'static str {
         match *self {
             Kind::Digraph => "digraph",
-            Kind::Graph => "graph"
+            Kind::Graph => "graph",
         }
     }
 
@@ -122,4 +139,4 @@ impl Kind {
             Kind::Graph => "--",
         }
     }
-}
+}