Add tests and refactor graph_from_edges

Author: prmtl

pydot_ng/__init__.py

@@ -14,6 +14,7 @@
 from __future__ import print_function
 
 import copy
+import functools
 import os
 import re
 import subprocess
@@ -306,27 +307,19 @@ def graph_from_edges(edge_list, node_prefix='', directed=False):
     If the graph is undirected by default, it is only
     calculated from one of the symmetric halves of the matrix.
     """
+    if edge_list is None:
+        edge_list = []
 
-    if directed:
-        graph = Dot(graph_type='digraph')
-
-    else:
-        graph = Dot(graph_type='graph')
+    graph_type = "digraph" if directed else "graph"
+    with_prefix = functools.partial("{0}{1}".format, node_prefix)
 
-    for edge in edge_list:
+    graph = Dot(graph_type=graph_type)
 
-        if isinstance(edge[0], str):
-            src = node_prefix + edge[0]
-        else:
-            src = node_prefix + str(edge[0])
-
-        if isinstance(edge[1], str):
-            dst = node_prefix + edge[1]
-        else:
-            dst = node_prefix + str(edge[1])
+    for src, dst in edge_list:
+        src = with_prefix(src)
+        dst = with_prefix(dst)
 
-        e = Edge(src, dst)
-        graph.add_edge(e)
+        graph.add_edge(Edge(src, dst))
 
     return graph
 
@@ -881,10 +874,14 @@ def get_source(self):
         """Get the edges source node name."""
         return self.obj_dict['points'][0]
 
+    source = property(get_source)
+
     def get_destination(self):
         """Get the edge's destination node name."""
         return self.obj_dict['points'][1]
 
+    destination = property(get_destination)
+
     def __hash__(self):
         return hash(hash(self.get_source()) + hash(self.get_destination()))
 

test/test_graph_from_edges.py

@@ -0,0 +1,90 @@
+# -*- coding: utf-8 -*-
+
+import functools
+from textwrap import dedent
+
+import pytest
+
+import pydot_ng
+
+
+@pytest.mark.parametrize(
+    "directed, graph_type",
+    ((True, "digraph"), (False, "graph")),
+    ids=("graph", "digraph"),
+)
+@pytest.mark.parametrize("edges", (None, [], ()), ids=(None, "list", "tuple"))
+def test_empty_list(directed, graph_type, edges):
+    graph = pydot_ng.graph_from_edges(edges, directed=directed)
+
+    assert graph.graph_type == graph_type
+    assert graph.to_string() == dedent(
+        """\
+        {graph_type} G {{
+        }}
+        """.format(
+            graph_type=graph_type
+        )
+    )
+    assert not graph.get_edges()
+
+
+@pytest.mark.parametrize(
+    "prefix", ("", "some_prefix_"), ids=("no prefix", "prefix")
+)
+@pytest.mark.parametrize("src", (1, 2.123, "a", "ą", True))
+@pytest.mark.parametrize("dst", (1, 2.123, "a", "ą", True))
+def test_edge_types(prefix, src, dst):
+    input_edges = [(src, dst)]
+    graph = pydot_ng.graph_from_edges(input_edges, node_prefix=prefix)
+
+    edges = graph.get_edges()
+    assert len(edges) == 1
+    edge = edges[0]
+
+    with_prefix = functools.partial("{0}{1}".format, prefix)
+
+    assert edge.source == pydot_ng.quote_if_necessary(with_prefix(src))
+    assert edge.destination == pydot_ng.quote_if_necessary(with_prefix(dst))
+
+
+@pytest.mark.parametrize(
+    "prefix, output",
+    (
+        (
+            "",
+            dedent(
+                """\
+        graph G {
+        1 -- 2;
+        2 -- "3.14";
+        "3.14" -- a;
+        a -- "ą";
+        "ą" -- True;
+        }
+        """
+            ),
+        ),
+        (
+            "prefix_",
+            dedent(
+                """\
+        graph G {
+        prefix_1 -- prefix_2;
+        prefix_2 -- "prefix_3.14";
+        "prefix_3.14" -- prefix_a;
+        prefix_a -- "prefix_ą";
+        "prefix_ą" -- prefix_True;
+        }
+        """
+            ),
+        ),
+    ),
+    ids=("no prefix", "prefix"),
+)
+def test_from_edge_to_string(prefix, output):
+    input_edges = [(1, 2), (2, 3.14), (3.14, "a"), ("a", "ą"), ("ą", True)]
+
+    graph = pydot_ng.graph_from_edges(input_edges, node_prefix=prefix)
+    assert len(graph.get_edges()) == len(input_edges)
+    assert graph.to_string() == output