Remove Attribute Statement Structs for IndexMap

I didnt think we were getting a ton of benefit from the
AttributeStatement abstraction so I replaced it with IndexMap which
the attribute statement impl were using internally.

While the DOT language docs (https://graphviz.org/doc/info/lang.html) do
call out attr_stmt as part of the language definition the overhead of
the related traits, impl, structs, etc felt a bit heavy. In particular,
supporting the ability to create them as part of a build as well as
allowing users to add to them via other bulid fns was a bit awkward.

For now I think removing the abstraction makes sense and provides for a
simpler implementation. Can revisit this down the road if other
requirements come up that perhaps warrant the addittional code.

3 files changed, 260 insertions, 291 deletions

diff --git a/src/attributes/mod.rs b/src/attributes/mod.rs
index 450440e..3219a50 100644
--- a/src/attributes/mod.rs
+++ b/src/attributes/mod.rs
@@ -321,24 +321,6 @@ pub enum AttributeType {
     Edge,
 }
 
-pub trait AttributeStatement<'a> {
-    fn get_attribute_statement_type(&self) -> &'static str;
-
-    fn get_attributes(&self) -> &IndexMap<String, AttributeText<'a>>;
-
-    fn dot_string(&self) -> String {
-        if self.get_attributes().is_empty() {
-            return String::from("");
-        }
-
-        format!(
-            "{}{};",
-            self.get_attribute_statement_type(),
-            fmt_attributes(self.get_attributes())
-        )
-    }
-}
-
 pub trait GraphAttributes<'a> {
     fn background(&mut self, background: String) -> &mut Self {
         self.add_attribute("_background", AttributeText::attr(background))
@@ -930,42 +912,8 @@ impl<'a> GraphAttributeStatementBuilder<'a> {
         }
     }
 
-    pub fn build(&self) -> GraphAttributeStatement<'a> {
-        GraphAttributeStatement {
-            attributes: self.attributes.clone(),
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct GraphAttributeStatement<'a> {
-    pub attributes: IndexMap<String, AttributeText<'a>>,
-}
-
-impl<'a> GraphAttributeStatement<'a> {
-    pub fn new() -> Self {
-        Self {
-            attributes: IndexMap::new(),
-        }
-    }
-
-    pub fn add_attribute<S: Into<String>>(
-        &mut self,
-        key: S,
-        value: AttributeText<'a>,
-    ) -> &mut Self {
-        self.attributes.insert(key.into(), value);
-        self
-    }
-}
-
-impl<'a> AttributeStatement<'a> for GraphAttributeStatement<'a> {
-    fn get_attribute_statement_type(&self) -> &'static str {
-        "graph"
-    }
-
-    fn get_attributes(&self) -> &IndexMap<String, AttributeText<'a>> {
-        &self.attributes
+    pub fn build(&self) -> IndexMap<String, AttributeText<'a>> {
+        self.attributes.clone()
     }
 }
 
@@ -1075,6 +1023,7 @@ impl Attributes {
         Self::add_attribute(attributes, "labelloc", AttributeText::from(label_location))
     }
 
+    // TODO: Layer, LayerRange, LayerList, LayerSep
     // TODO: layer struct
     pub fn layer(attributes: &mut IndexMap<String, AttributeText>, layer: String) {
         Self::add_attribute(attributes, "layer", AttributeText::attr(layer))
@@ -1233,8 +1182,8 @@ pub trait NodeAttributes<'a> {
     /// If true, the node size is specified by the values of the width and height attributes only and
     /// is not expanded to contain the text label.
     /// There will be a warning if the label (with margin) cannot fit within these limits.
-    /// If false, the size of a node is determined by smallest width and height needed to contain its label
-    /// and image, if any, with a margin specified by the margin attribute.
+    /// If false, the size of a node is determined by smallest width and height needed
+    /// to contain its label and image, if any, with a margin specified by the margin attribute.
     fn fixed_size(&mut self, fixed_size: bool) -> &mut Self {
         self.add_attribute("fixedsize", AttributeText::from(fixed_size))
     }
@@ -1307,11 +1256,15 @@ pub trait NodeAttributes<'a> {
     }
 
     // Vertical placement of labels for nodes, root graphs and clusters.
-    // For graphs and clusters, only labelloc=t and labelloc=b are allowed, corresponding to placement at the top and bottom, respectively.
+    // For graphs and clusters, only labelloc=t and labelloc=b are allowed,
+    // corresponding to placement at the top and bottom, respectively.
     // By default, root graph labels go on the bottom and cluster labels go on the top.
-    // Note that a subgraph inherits attributes from its parent. Thus, if the root graph sets labelloc=b, the subgraph inherits this value.
-    // For nodes, this attribute is used only when the height of the node is larger than the height of its label.
-    // If labelloc=t, labelloc=c, labelloc=b, the label is aligned with the top, centered, or aligned with the bottom of the node, respectively.
+    // Note that a subgraph inherits attributes from its parent.
+    // Thus, if the root graph sets labelloc=b, the subgraph inherits this value.
+    // For nodes, this attribute is used only when the height of the node is
+    // larger than the height of its label.
+    // If labelloc=t, labelloc=c, labelloc=b, the label is aligned with the top,
+    // centered, or aligned with the bottom of the node, respectively.
     // By default, the label is vertically centered.
     fn label_location(&mut self, label_location: LabelLocation) -> &mut Self {
         Attributes::label_location(self.get_attributes_mut(), label_location);
@@ -1346,12 +1299,17 @@ pub trait NodeAttributes<'a> {
         self
     }
 
-    /// By default, the justification of multi-line labels is done within the largest context that makes sense.
-    /// Thus, in the label of a polygonal node, a left-justified line will align with the left side of the node (shifted by the prescribed margin).
-    /// In record nodes, left-justified line will line up with the left side of the enclosing column of fields.
+    /// By default, the justification of multi-line labels is done within the
+    /// largest context that makes sense.
+    /// Thus, in the label of a polygonal node, a left-justified line will align
+    /// with the left side of the node (shifted by the prescribed margin).
+    /// In record nodes, left-justified line will line up with the left side of
+    /// the enclosing column of fields.
     /// If nojustify=true, multi-line labels will be justified in the context of itself.
-    /// For example, if nojustify is set, the first label line is long, and the second is shorter and left-justified,
-    /// the second will align with the left-most character in the first line, regardless of how large the node might be.
+    /// For example, if nojustify is set, the first label line is long, and the
+    /// second is shorter and left-justified,
+    /// the second will align with the left-most character in the first line,
+    /// regardless of how large the node might be.
     fn no_justify(&mut self, no_justify: bool) -> &mut Self {
         Attributes::no_justify(self.get_attributes_mut(), no_justify);
         self
@@ -1426,7 +1384,8 @@ pub trait NodeAttributes<'a> {
     }
 
     // TODO: constrain
-    /// Print guide boxes in PostScript at the beginning of routesplines if showboxes=1, or at the end if showboxes=2.
+    /// Print guide boxes in PostScript at the beginning of routesplines if
+    /// showboxes=1, or at the end if showboxes=2.
     /// (Debugging, TB mode only!)
     /// default: 0, minimum: 0
     fn show_boxes(&mut self, show_boxes: u32) -> &mut Self {
@@ -1461,7 +1420,8 @@ pub trait NodeAttributes<'a> {
         self
     }
 
-    /// If the object has a URL, this attribute determines which window of the browser is used for the URL.
+    /// If the object has a URL, this attribute determines which window of the
+    /// browser is used for the URL.
     fn target(&mut self, target: String) -> &mut Self {
         Attributes::target(self.get_attributes_mut(), target);
         self
@@ -1471,7 +1431,8 @@ pub trait NodeAttributes<'a> {
     /// If unset, Graphviz will use the object’s label if defined.
     /// Note that if the label is a record specification or an HTML-like label,
     /// the resulting tooltip may be unhelpful.
-    /// In this case, if tooltips will be generated, the user should set a tooltip attribute explicitly.
+    /// In this case, if tooltips will be generated, the user should set a
+    /// tooltip attribute explicitly.
     fn tooltip(&mut self, tooltip: String) -> &mut Self {
         Attributes::tooltip(self.get_attributes_mut(), tooltip);
         self
@@ -1609,7 +1570,8 @@ pub trait EdgeAttributes<'a> {
 
     /// Edge type for drawing arrowheads.
     /// Indicates which ends of the edge should be decorated with an arrowhead.
-    /// The actual style of the arrowhead can be specified using the arrowhead and arrowtail attributes.
+    /// The actual style of the arrowhead can be specified using the arrowhead
+    /// and arrowtail attributes.
     fn dir(&mut self, dir: Direction) -> &mut Self {
         self.add_attribute("dir", AttributeText::from(dir))
     }
@@ -1630,19 +1592,30 @@ pub trait EdgeAttributes<'a> {
 
     /// The link for the non-label parts of an edge.
     /// edgeURL overrides any URL defined for the edge.
-    /// Also, edgeURL is used near the head or tail node unless overridden by headURL or tailURL, respectively.
+    /// Also, edgeURL is used near the head or tail node unless overridden by
+    /// headURL or tailURL, respectively.
     fn edge_url(&mut self, edge_url: String) -> &mut Self {
         self.add_attribute("edgeurl", AttributeText::escaped(edge_url))
     }
 
-    // TODO: color list
-    /// Color used to fill the background of a node or cluster assuming style=filled, or a filled arrowhead.
+    /// Color used to fill the background of a node or cluster assuming
+    /// style=filled, or a filled arrowhead.
     fn fill_color(&mut self, fill_color: Color<'a>) -> &mut Self {
         Attributes::fill_color(self.get_attributes_mut(), fill_color);
         self
     }
 
-    // TODO: color list
+    /// Color used to fill the background of a node or cluster assuming
+    /// style=filled, or a filled arrowhead.
+    /// A gradient fill is used. By default, this is a linear fill; setting style=radial
+    /// will cause a radial fill.
+    /// At present, only two colors are used. If the second color (after a colon) is missing,
+    /// the default color is used for it.
+    fn fill_color_with_colorlist(&mut self, fill_color: ColorList<'a>) -> &mut Self {
+        Attributes::fill_color_with_colorlist(self.get_attributes_mut(), fill_color);
+        self
+    }
+
     /// Color used for text.
     fn font_color(&mut self, font_color: Color<'a>) -> &mut Self {
         Attributes::font_color(self.get_attributes_mut(), font_color);
@@ -1936,20 +1909,6 @@ pub trait EdgeAttributes<'a> {
     ) -> &mut Self;
 
     fn get_attributes_mut(&mut self) -> &mut IndexMap<String, AttributeText<'a>>;
-
-    // fn add_attribute<S: Into<String>>(
-    //     &self,
-    //     key: S,
-    //     value: AttributeText<'a>
-    // ) {
-    //     self.get_attributes().insert(key.into(), value);
-    // }
-
-    // fn get_attributes(&self) -> IndexMap<String, AttributeText<'a>>;
-
-    // fn get_attributes_mut(&self) -> &mut IndexMap<String, AttributeText<'a>>;
-
-    // fn to_dot_string(&self) -> String;
 }
 
 pub(crate) fn fmt_attributes(attributes: &IndexMap<String, AttributeText>) -> String {
@@ -1972,8 +1931,10 @@ pub(crate) fn fmt_attributes(attributes: &IndexMap<String, AttributeText>) -> St
 #[cfg(test)]
 mod test {
     use crate::attributes::{
-        AttributeStatement, Color, GraphAttributeStatementBuilder, GraphAttributes,
+        fmt_attributes, AttributeText, Color, GraphAttributeStatementBuilder,
+        GraphAttributes,
     };
+    use indexmap::map::IndexMap;
 
     #[test]
     fn graph_attribute_colorlist_vec_dot_string() {
@@ -1985,8 +1946,30 @@ mod test {
             .build();
 
         assert_eq!(
-            "graph [fillcolor=\"yellow;0.3:blue\"];",
-            graph_attributes.dot_string()
+            graph_attributes.get("fillcolor").unwrap().dot_string(),
+            "\"yellow;0.3:blue\""
         );
     }
+
+    #[test]
+    fn fmt_attributes_empty_attributes_should_return_empty_string() {
+        assert_eq!(fmt_attributes(&IndexMap::new()), "");
+    }
+
+    #[test]
+    fn fmt_attributes_with_single_attribute() {
+        let mut attributes = IndexMap::new();
+        attributes.insert("color".to_string(), AttributeText::attr("red"));
+
+        assert_eq!(fmt_attributes(&attributes), " [color=red]");
+    }
+
+    #[test]
+    fn fmt_attributes_with_attributes() {
+        let mut attributes = IndexMap::new();
+        attributes.insert("color".to_string(), AttributeText::attr("red"));
+        attributes.insert("size".to_string(), AttributeText::attr("2"));
+
+        assert_eq!(fmt_attributes(&attributes), " [color=red, size=2]");
+    }
 }
diff --git a/src/dot.rs b/src/dot.rs
index 1cbb175..5c1dc29 100644
--- a/src/dot.rs
+++ b/src/dot.rs
@@ -1,8 +1,8 @@
 // TODO: docs
 
 use crate::attributes::{
-    fmt_attributes, AttributeStatement, AttributeText, AttributeType, EdgeAttributes,
-    GraphAttributeStatement, NodeAttributes, PortPosition,
+    fmt_attributes, AttributeText, AttributeType, EdgeAttributes, NodeAttributes,
+    PortPosition,
 };
 use indexmap::IndexMap;
 use std::borrow::Cow;
@@ -49,16 +49,31 @@ impl<'a> Dot<'a> {
 
         writeln!(w, " {{")?;
 
-        if let Some(graph_attributes) = &graph.graph_attributes {
-            write!(w, "{}{}\n", INDENT, graph_attributes.dot_string())?;
+        if !&graph.graph_attributes.is_empty() {
+            write!(
+                w,
+                "{}graph{};\n",
+                INDENT,
+                fmt_attributes(&graph.graph_attributes)
+            )?;
         }
 
-        if let Some(node_attributes) = &graph.node_attributes {
-            write!(w, "{}{}\n", INDENT, node_attributes.dot_string())?;
+        if !&graph.node_attributes.is_empty() {
+            write!(
+                w,
+                "{}node{};\n",
+                INDENT,
+                fmt_attributes(&graph.node_attributes)
+            )?;
         }
 
-        if let Some(edge_attributes) = &graph.edge_attributes {
-            write!(w, "{}{}\n", INDENT, edge_attributes.dot_string())?;
+        if !&graph.edge_attributes.is_empty() {
+            write!(
+                w,
+                "{}edge{};\n",
+                INDENT,
+                fmt_attributes(&graph.edge_attributes)
+            )?;
         }
 
         for g in &graph.sub_graphs {
@@ -94,16 +109,32 @@ impl<'a> Dot<'a> {
         writeln!(w, " {{")?;
 
         let indent = get_indentation(indentation_level + 1);
-        if let Some(graph_attributes) = &sub_graph.graph_attributes {
-            write!(w, "{}{}\n", indent, graph_attributes.dot_string())?;
+
+        if !&sub_graph.graph_attributes.is_empty() {
+            write!(
+                w,
+                "{}graph{};\n",
+                indent,
+                fmt_attributes(&sub_graph.graph_attributes)
+            )?;
         }
 
-        if let Some(node_attributes) = &sub_graph.node_attributes {
-            write!(w, "{}{}\n", indent, node_attributes.dot_string())?;
+        if !&sub_graph.node_attributes.is_empty() {
+            write!(
+                w,
+                "{}node{};\n",
+                indent,
+                fmt_attributes(&sub_graph.node_attributes)
+            )?;
         }
 
-        if let Some(edge_attributes) = &sub_graph.edge_attributes {
-            write!(w, "{}{}\n", indent, edge_attributes.dot_string())?;
+        if !&sub_graph.edge_attributes.is_empty() {
+            write!(
+                w,
+                "{}edge{};\n",
+                indent,
+                fmt_attributes(&sub_graph.edge_attributes)
+            )?;
         }
 
         for g in &sub_graph.sub_graphs {
@@ -193,11 +224,11 @@ pub struct Graph<'a> {
     /// Comment added to the first line of the source.
     pub comment: Option<String>,
 
-    pub graph_attributes: Option<GraphAttributeStatement<'a>>,
+    pub graph_attributes: IndexMap<String, AttributeText<'a>>,
 
-    pub node_attributes: Option<NodeAttributeStatement<'a>>,
+    pub node_attributes: IndexMap<String, AttributeText<'a>>,
 
-    pub edge_attributes: Option<EdgeAttributeStatement<'a>>,
+    pub edge_attributes: IndexMap<String, AttributeText<'a>>,
 
     pub sub_graphs: Vec<SubGraph<'a>>,
 
@@ -212,9 +243,9 @@ impl<'a> Graph<'a> {
         is_directed: bool,
         strict: bool,
         comment: Option<String>,
-        graph_attributes: Option<GraphAttributeStatement<'a>>,
-        node_attributes: Option<NodeAttributeStatement<'a>>,
-        edge_attributes: Option<EdgeAttributeStatement<'a>>,
+        graph_attributes: IndexMap<String, AttributeText<'a>>,
+        node_attributes: IndexMap<String, AttributeText<'a>>,
+        edge_attributes: IndexMap<String, AttributeText<'a>>,
         sub_graphs: Vec<SubGraph<'a>>,
         nodes: Vec<Node<'a>>,
         edges: Vec<Edge<'a>>,
@@ -257,11 +288,11 @@ pub struct GraphBuilder<'a> {
 
     strict: bool,
 
-    graph_attributes: Option<GraphAttributeStatement<'a>>,
+    graph_attributes: IndexMap<String, AttributeText<'a>>,
 
-    node_attributes: Option<NodeAttributeStatement<'a>>,
+    node_attributes: IndexMap<String, AttributeText<'a>>,
 
-    edge_attributes: Option<EdgeAttributeStatement<'a>>,
+    edge_attributes: IndexMap<String, AttributeText<'a>>,
 
     sub_graphs: Vec<SubGraph<'a>>,
 
@@ -279,9 +310,9 @@ impl<'a> GraphBuilder<'a> {
             id,
             is_directed: true,
             strict: false,
-            graph_attributes: None,
-            node_attributes: None,
-            edge_attributes: None,
+            graph_attributes: IndexMap::new(),
+            node_attributes: IndexMap::new(),
+            edge_attributes: IndexMap::new(),
             sub_graphs: Vec::new(),
             nodes: Vec::new(),
             edges: Vec::new(),
@@ -294,9 +325,9 @@ impl<'a> GraphBuilder<'a> {
             id,
             is_directed: false,
             strict: false,
-            graph_attributes: None,
-            node_attributes: None,
-            edge_attributes: None,
+            graph_attributes: IndexMap::new(),
+            node_attributes: IndexMap::new(),
+            edge_attributes: IndexMap::new(),
             sub_graphs: Vec::new(),
             nodes: Vec::new(),
             edges: Vec::new(),
@@ -311,47 +342,28 @@ impl<'a> GraphBuilder<'a> {
 
     pub fn add_graph_attributes(
         &mut self,
-        graph_attributes: GraphAttributeStatement<'a>,
+        attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.graph_attributes = Some(graph_attributes);
+        self.graph_attributes.extend(attributes);
         self
     }
 
     pub fn add_node_attributes(
         &mut self,
-        node_attributes: NodeAttributeStatement<'a>,
+        node_attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.node_attributes = Some(node_attributes);
+        self.node_attributes.extend(node_attributes);
         self
     }
 
     pub fn add_edge_attributes(
         &mut self,
-        edge_attributes: EdgeAttributeStatement<'a>,
+        edge_attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.edge_attributes = Some(edge_attributes);
+        self.edge_attributes.extend(edge_attributes);
         self
     }
 
-    // TODO: update to insert into appropriate statement or remove?
-    // pub fn add_attribute(
-    //     &mut self,
-    //     attribute_type: AttributeType,
-    //     key: String, value: AttributeText<'a>
-    // ) -> &mut Self {
-    //     self.get_attributes(attribute_type).insert(key, value);
-    //     self
-    // }
-    //
-    // pub fn add_attributes(
-    //     &mut self,
-    //     attribute_type: AttributeType,
-    //     attributes: HashMap<String, AttributeText<'a>>
-    // ) -> &mut Self {
-    //     self.get_attributes(attribute_type).extend(attributes);
-    //     self
-    // }
-
     pub fn add_attribute(
         &mut self,
         attribute_type: AttributeType,
@@ -359,34 +371,23 @@ impl<'a> GraphBuilder<'a> {
         value: AttributeText<'a>,
     ) -> &mut Self {
         match attribute_type {
-            AttributeType::Graph => {
-                if self.graph_attributes.is_none() {
-                    self.graph_attributes = Some(GraphAttributeStatement::new());
-                }
-                self.graph_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
-            }
-            AttributeType::Edge => {
-                if self.edge_attributes.is_none() {
-                    self.edge_attributes = Some(EdgeAttributeStatement::new());
-                }
-                self.edge_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
-            }
-            AttributeType::Node => {
-                if self.node_attributes.is_none() {
-                    self.node_attributes = Some(NodeAttributeStatement::new());
-                }
-                self.node_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
-            }
-        }
+            AttributeType::Graph => self.graph_attributes.insert(key, value),
+            AttributeType::Edge => self.edge_attributes.insert(key, value),
+            AttributeType::Node => self.node_attributes.insert(key, value),
+        };
+        self
+    }
+
+    pub fn extend_with_attributes(
+        &mut self,
+        attribute_type: AttributeType,
+        attributes: HashMap<String, AttributeText<'a>>,
+    ) -> &mut Self {
+        match attribute_type {
+            AttributeType::Graph => self.graph_attributes.extend(attributes),
+            AttributeType::Edge => self.edge_attributes.extend(attributes),
+            AttributeType::Node => self.node_attributes.extend(attributes),
+        };
         self
     }
 
@@ -430,12 +431,11 @@ impl<'a> GraphBuilder<'a> {
 pub struct SubGraph<'a> {
     pub id: Option<String>,
 
-    pub graph_attributes: Option<GraphAttributeStatement<'a>>,
-
-    pub node_attributes: Option<NodeAttributeStatement<'a>>,
+    pub graph_attributes: IndexMap<String, AttributeText<'a>>,
 
-    pub edge_attributes: Option<EdgeAttributeStatement<'a>>,
+    pub node_attributes: IndexMap<String, AttributeText<'a>>,
 
+    pub edge_attributes: IndexMap<String, AttributeText<'a>>,
     pub sub_graphs: Vec<SubGraph<'a>>,
 
     pub nodes: Vec<Node<'a>>,
@@ -446,9 +446,9 @@ pub struct SubGraph<'a> {
 impl<'a> SubGraph<'a> {
     pub fn new(
         id: Option<String>,
-        graph_attributes: Option<GraphAttributeStatement<'a>>,
-        node_attributes: Option<NodeAttributeStatement<'a>>,
-        edge_attributes: Option<EdgeAttributeStatement<'a>>,
+        graph_attributes: IndexMap<String, AttributeText<'a>>,
+        node_attributes: IndexMap<String, AttributeText<'a>>,
+        edge_attributes: IndexMap<String, AttributeText<'a>>,
         sub_graphs: Vec<SubGraph<'a>>,
         nodes: Vec<Node<'a>>,
         edges: Vec<Edge<'a>>,
@@ -468,11 +468,11 @@ impl<'a> SubGraph<'a> {
 pub struct SubGraphBuilder<'a> {
     id: Option<String>,
 
-    graph_attributes: Option<GraphAttributeStatement<'a>>,
+    graph_attributes: IndexMap<String, AttributeText<'a>>,
 
-    node_attributes: Option<NodeAttributeStatement<'a>>,
+    node_attributes: IndexMap<String, AttributeText<'a>>,
 
-    edge_attributes: Option<EdgeAttributeStatement<'a>>,
+    edge_attributes: IndexMap<String, AttributeText<'a>>,
 
     sub_graphs: Vec<SubGraph<'a>>,
 
@@ -486,9 +486,9 @@ impl<'a> SubGraphBuilder<'a> {
     pub fn new(id: Option<String>) -> Self {
         Self {
             id,
-            graph_attributes: None,
-            node_attributes: None,
-            edge_attributes: None,
+            graph_attributes: IndexMap::new(),
+            node_attributes: IndexMap::new(),
+            edge_attributes: IndexMap::new(),
             sub_graphs: Vec::new(),
             nodes: Vec::new(),
             edges: Vec::new(),
@@ -497,25 +497,25 @@ impl<'a> SubGraphBuilder<'a> {
 
     pub fn add_graph_attributes(
         &mut self,
-        graph_attributes: GraphAttributeStatement<'a>,
+        graph_attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.graph_attributes = Some(graph_attributes);
+        self.graph_attributes.extend(graph_attributes);
         self
     }
 
     pub fn add_node_attributes(
         &mut self,
-        node_attributes: NodeAttributeStatement<'a>,
+        node_attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.node_attributes = Some(node_attributes);
+        self.node_attributes.extend(node_attributes);
         self
     }
 
     pub fn add_edge_attributes(
         &mut self,
-        edge_attributes: EdgeAttributeStatement<'a>,
+        edge_attributes: IndexMap<String, AttributeText<'a>>,
     ) -> &mut Self {
-        self.edge_attributes = Some(edge_attributes);
+        self.edge_attributes.extend(edge_attributes);
         self
     }
 
@@ -546,31 +546,13 @@ impl<'a> SubGraphBuilder<'a> {
     ) -> &mut Self {
         match attribute_type {
             AttributeType::Graph => {
-                if self.graph_attributes.is_none() {
-                    self.graph_attributes = Some(GraphAttributeStatement::new());
-                }
-                self.graph_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
+                self.graph_attributes.insert(key, value);
             }
             AttributeType::Edge => {
-                if self.edge_attributes.is_none() {
-                    self.edge_attributes = Some(EdgeAttributeStatement::new());
-                }
-                self.edge_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
+                self.edge_attributes.insert(key, value);
             }
             AttributeType::Node => {
-                if self.node_attributes.is_none() {
-                    self.node_attributes = Some(NodeAttributeStatement::new());
-                }
-                self.node_attributes
-                    .as_mut()
-                    .unwrap()
-                    .add_attribute(key, value);
+                self.node_attributes.insert(key, value);
             }
         }
         self
@@ -732,12 +714,6 @@ impl<'a> EdgeAttributes<'a> for EdgeBuilder<'a> {
     fn get_attributes_mut(&mut self) -> &mut IndexMap<String, AttributeText<'a>> {
         &mut self.attributes
     }
-
-    // /// Add multiple attributes to the edge.
-    // fn add_attributes(&'a mut self, attributes: HashMap<String, AttributeText<'a>>) -> &mut Self {
-    //     self.attributes.extend(attributes);
-    //     self
-    // }
 }
 
 impl<'a> EdgeBuilder<'a> {
@@ -842,42 +818,8 @@ impl<'a> NodeAttributeStatementBuilder<'a> {
         }
     }
 
-    pub fn build(&self) -> NodeAttributeStatement<'a> {
-        NodeAttributeStatement {
-            attributes: self.attributes.clone(),
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct NodeAttributeStatement<'a> {
-    pub attributes: IndexMap<String, AttributeText<'a>>,
-}
-
-impl<'a> NodeAttributeStatement<'a> {
-    pub fn new() -> Self {
-        Self {
-            attributes: IndexMap::new(),
-        }
-    }
-
-    pub fn add_attribute<S: Into<String>>(
-        &mut self,
-        key: S,
-        value: AttributeText<'a>,
-    ) -> &mut Self {
-        self.attributes.insert(key.into(), value);
-        self
-    }
-}
-
-impl<'a> AttributeStatement<'a> for NodeAttributeStatement<'a> {
-    fn get_attribute_statement_type(&self) -> &'static str {
-        "node"
-    }
-
-    fn get_attributes(&self) -> &IndexMap<String, AttributeText<'a>> {
-        &self.attributes
+    pub fn build(&self) -> IndexMap<String, AttributeText<'a>> {
+        self.attributes.clone()
     }
 }
 
@@ -896,16 +838,6 @@ impl<'a> EdgeAttributes<'a> for EdgeAttributeStatementBuilder<'a> {
     }
 }
 
-impl<'a> AttributeStatement<'a> for EdgeAttributeStatement<'a> {
-    fn get_attribute_statement_type(&self) -> &'static str {
-        "edge"
-    }
-
-    fn get_attributes(&self) -> &IndexMap<String, AttributeText<'a>> {
-        &self.attributes
-    }
-}
-
 // I'm not a huge fan of needing this builder but having a hard time getting around &mut without it
 pub struct EdgeAttributeStatementBuilder<'a> {
     pub attributes: IndexMap<String, AttributeText<'a>>,
@@ -918,32 +850,8 @@ impl<'a> EdgeAttributeStatementBuilder<'a> {
         }
     }
 
-    pub fn build(&self) -> EdgeAttributeStatement<'a> {
-        EdgeAttributeStatement {
-            attributes: self.attributes.clone(),
-        }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct EdgeAttributeStatement<'a> {
-    pub attributes: IndexMap<String, AttributeText<'a>>,
-}
-
-impl<'a> EdgeAttributeStatement<'a> {
-    pub fn new() -> Self {
-        Self {
-            attributes: IndexMap::new(),
-        }
-    }
-
-    pub fn add_attribute<S: Into<String>>(
-        &mut self,
-        key: S,
-        value: AttributeText<'a>,
-    ) -> &mut Self {
-        self.attributes.insert(key.into(), value);
-        self
+    pub fn build(&self) -> IndexMap<String, AttributeText<'a>> {
+        self.attributes.clone()
     }
 }
 
diff --git a/tests/dot.rs b/tests/dot.rs
index 526ff83..c4d31ba 100644
--- a/tests/dot.rs
+++ b/tests/dot.rs
@@ -1,5 +1,5 @@
 use dotavious::attributes::{
-    AttributeText, Color, CompassPoint, EdgeAttributes, EdgeStyle,
+    AttributeText, AttributeType, Color, CompassPoint, EdgeAttributes, EdgeStyle,
     GraphAttributeStatementBuilder, GraphAttributes, GraphStyle, NodeAttributes,
     NodeStyle, PortPosition, RankDir, Shape,
 };
@@ -298,6 +298,84 @@ fn port_position_attribute() {
 
 #[test]
 fn graph_attributes() {
+    let g = GraphBuilder::new_directed(Some("graph_attributes".to_string()))
+        .add_attribute(
+            AttributeType::Graph,
+            "rankdir".to_string(),
+            AttributeText::from(RankDir::LeftRight),
+        )
+        .add_attribute(
+            AttributeType::Node,
+            "style".to_string(),
+            AttributeText::from(NodeStyle::Filled),
+        )
+        .add_attribute(
+            AttributeType::Edge,
+            "color".to_string(),
+            AttributeText::from(Color::Named("red")),
+        )
+        .build();
+
+    let r = test_input(g);
+
+    assert_eq!(
+        r.unwrap(),
+        r#"digraph graph_attributes {
+    graph [rankdir=LR];
+    node [style=filled];
+    edge [color="red"];
+}
+"#
+    );
+}
+
+#[test]
+fn graph_attributes_extend() {
+    let g = GraphBuilder::new_directed(Some("graph_attributes".to_string()))
+        .extend_with_attributes(
+            AttributeType::Graph,
+            [(
+                "rankdir".to_string(),
+                AttributeText::from(RankDir::LeftRight),
+            )]
+            .iter()
+            .cloned()
+            .collect(),
+        )
+        .extend_with_attributes(
+            AttributeType::Node,
+            [("style".to_string(), AttributeText::from(NodeStyle::Filled))]
+                .iter()
+                .cloned()
+                .collect(),
+        )
+        .extend_with_attributes(
+            AttributeType::Edge,
+            [(
+                "color".to_string(),
+                AttributeText::from(Color::Named("red")),
+            )]
+            .iter()
+            .cloned()
+            .collect(),
+        )
+        .build();
+
+    let r = test_input(g);
+
+    assert_eq!(
+        r.unwrap(),
+        r#"digraph graph_attributes {
+    graph [rankdir=LR];
+    node [style=filled];
+    edge [color="red"];
+}
+"#
+    );
+}
+
+#[test]
+fn graph_attributes_statement_builders() {
     let graph_attributes = GraphAttributeStatementBuilder::new()
         .rank_dir(RankDir::LeftRight)
         .build();