diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 6968bc4a..2a388f5f 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -83,6 +83,7 @@ class Diagram: direction: str = "LR", curvestyle: str = "ortho", outformat: str = "png", + autolabel: bool = False, show: bool = True, graph_attr: dict = {}, node_attr: dict = {}, @@ -137,6 +138,7 @@ class Diagram: self.dot.edge_attr.update(edge_attr) self.show = show + self.autolabel = autolabel def __str__(self) -> str: return str(self.dot) @@ -296,11 +298,20 @@ class Node: self._id = self._rand_id() self.label = label + # Node must be belong to a diagrams. + self._diagram = getdiagram() + if self._diagram is None: + raise EnvironmentError("Global diagrams context not set up") + + if self._diagram.autolabel: + # import pdb; pdb.set_trace() + self.label = self.__class__.__name__ + "\n" + self.label + # fmt: off # If a node has an icon, increase the height slightly to avoid # that label being spanned between icon image and white space. # Increase the height by the number of new lines included in the label. - padding = 0.4 * (label.count('\n')) + padding = 0.4 * (self.label.count('\n')) self._attrs = { "shape": "none", "height": str(self._height + padding), @@ -310,10 +321,6 @@ class Node: # fmt: on self._attrs.update(attrs) - # Node must be belong to a diagrams. - self._diagram = getdiagram() - if self._diagram is None: - raise EnvironmentError("Global diagrams context not set up") self._cluster = getcluster() # If a node is in the cluster context, add it to cluster. diff --git a/tests/test_diagram.py b/tests/test_diagram.py index ad8558c5..f4efa1e3 100644 --- a/tests/test_diagram.py +++ b/tests/test_diagram.py @@ -106,6 +106,12 @@ class DiagramTest(unittest.TestCase): with Diagram(show=False): Node("node1") self.assertTrue(os.path.exists(f"{self.name}.png")) + + def test_autolabel(self): + with Diagram(name=os.path.join(self.name, "nodes_to_node"), show=False): + node1 = Node("node1") + self.assertTrue(node1.label,"Node\nnode1") + class ClusterTest(unittest.TestCase):