Replace invalid symbols in the labels for metadata visualization

sdv/metadata/visualization.py

@@ -70,8 +70,12 @@ def _get_graphviz_extension(filepath):
     return None, None
 
 
+def _replace_special_characters(string):
+    return string.replace('<', '\<').replace('>', '\>')  # noqa: W605
+
+
 def visualize_graph(nodes, edges, filepath=None):
-    """Plot metadata usign graphviz.
+    """Plot metadata using graphviz.
 
     Try to generate a plot using graphviz.
     If a ``filepath`` is provided save the output into a file.
@@ -105,10 +109,10 @@ def visualize_graph(nodes, edges, filepath=None):
     )
 
     for name, label in nodes.items():
-        digraph.node(name, label=label)
+        digraph.node(name, label=_replace_special_characters(label))
 
     for parent, child, label in edges:
-        digraph.edge(parent, child, label=label, arrowhead='oinv')
+        digraph.edge(parent, child, label=_replace_special_characters(label), arrowhead='oinv')
 
     if filename:
         digraph.render(filename=filename, cleanup=True, format=graphviz_extension)

tests/integration/metadata/test_visualization.py

@@ -0,0 +1,31 @@
+import pandas as pd
+
+from sdv.metadata import MultiTableMetadata, SingleTableMetadata
+
+
+def test_visualize_graph_for_single_table():
+    """Test it runs when a column name contains `>`."""
+    # Setup
+    data = pd.DataFrame({'>': ['a', 'b', 'c']})
+    metadata = SingleTableMetadata()
+    metadata.detect_from_dataframe(data)
+
+    # Run
+    metadata.visualize()
+
+
+def test_visualize_graph_for_multi_table():
+    """Test it runs when a column name contains `>`."""
+    # Setup
+    data1 = pd.DataFrame({'>': ['a', 'b', 'c']})
+    data2 = pd.DataFrame({'>': ['a', 'b', 'c']})
+    tables = {'1': data1, '2': data2}
+    metadata = MultiTableMetadata()
+    metadata.detect_from_dataframes(tables)
+    metadata.update_column('1', '>', sdtype='id')
+    metadata.update_column('2', '>', sdtype='id')
+    metadata.set_primary_key('1', '>')
+    metadata.add_relationship('1', '2', '>', '>')
+
+    # Run
+    metadata.visualize()