Correctly return None for label for invalid edge/vertex and default to empty string when no label has been set

Author: matt-snider

src/dot.rs

@@ -48,17 +48,13 @@ impl<'a, T> dot::Labeller<'a, Nd, Ed<'a>> for DotGraph<'a, T> {
     }
 
     fn node_label<'b>(&'b self, n: &Nd) -> dot::LabelText<'b> {
-        let label = self.graph.vertex_label(n)
-            .cloned()
-            .unwrap_or_else(|| String::new());
-        dot::LabelText::label(Cow::Owned(label))
+        let label = self.graph.vertex_label(n).unwrap();
+        dot::LabelText::label(Cow::Borrowed(label))
     }
 
     fn edge_label<'b>(&'b self, e: &Ed) -> dot::LabelText<'b> {
-        let label = self.graph.edge_label(e.0, e.1)
-            .cloned()
-            .unwrap_or_else(|| String::new());
-        dot::LabelText::LabelStr(Cow::Owned(label))
+        let label = self.graph.edge_label(e.0, e.1).unwrap();
+        dot::LabelText::LabelStr(Cow::Borrowed(label))
     }
 }
 

src/graph.rs

@@ -17,8 +17,6 @@ use std::fmt::Debug;
 
 #[cfg(feature = "no_std")]
 use core::mem;
-#[cfg(not(feature = "no_std"))]
-use std::mem;
 
 #[cfg(feature = "no_std")]
 extern crate alloc;
@@ -91,6 +89,8 @@ pub struct Graph<T> {
     edge_labels: HashMap<Edge, String>,
 }
 
+const DEFAULT_LABEL: &str = "";
+
 impl<T> Graph<T> {
     /// Creates a new graph.
     ///
@@ -1437,8 +1437,14 @@ impl<T> Graph<T> {
     /// This method requires the `dot` crate feature.
     ///
     /// Returns `None` if there is no vertex associated with the given id in the graph.
-    pub fn vertex_label(&self, vertex_id: &VertexId) -> Option<&String> {
+    pub fn vertex_label(&self, vertex_id: &VertexId) -> Option<&str> {
+        if !self.vertices.contains_key(vertex_id) {
+            return None;
+        }
+
         self.vertex_labels.get(vertex_id)
+            .map(|x| x.as_str())
+            .or(Some(&DEFAULT_LABEL))
     }
 
     #[cfg(feature = "dot")]
@@ -1447,8 +1453,14 @@ impl<T> Graph<T> {
     /// This method requires the `dot` crate feature.
     ///
     /// Returns `None` if there is no edge associated with the given vertices in the graph.
-    pub fn edge_label(&self, a: &VertexId, b: &VertexId) -> Option<&String> {
+    pub fn edge_label(&self, a: &VertexId, b: &VertexId) -> Option<&str> {
+        if !self.has_edge(a, b) {
+            return None;
+        }
+
         self.edge_labels.get(&Edge::new(*a, *b))
+            .map(|x| x.as_str())
+            .or(Some(&DEFAULT_LABEL))
     }
 
     #[cfg(feature = "dot")]