From 55d3e73184b2ee0f3ed33da40750e0a5a3c8defa Mon Sep 17 00:00:00 2001
From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com>
Date: Tue, 29 Dec 2020 15:25:55 -0300
Subject: [PATCH] Convert the "Cluster" class to a "Node" and extract commons
 to a base class

---
 diagrams/__init__.py | 79 ++++++++++++++++++++++----------------------
 1 file changed, 39 insertions(+), 40 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index 281ec09f..1d4c3207 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -17,7 +17,10 @@ __cluster = contextvars.ContextVar("cluster")
 
 
 def getdiagram():
-    return __diagram.get()
+    try:
+        return __diagram.get()
+    except LookupError:
+        raise EnvironmentError("Global diagrams context not set up")
 
 
 def setdiagram(diagram):
@@ -50,7 +53,7 @@ class _Cluster:
 
         try:
             self._parent = getcluster() or getdiagram()
-        except LookupError:
+        except EnvironmentError:
             self._parent = None
 
 
@@ -58,40 +61,28 @@ class _Cluster:
         setcluster(self)
         return self
 
-    def __exit__(self, *args):
+    def __exit__(self, exc_type, exc_value, traceback):
         setcluster(self._parent)
 
-        if not (self.nodes or self.subgraphs):
-            return
+        for nodeid, node in self.nodes.items():
+            self.dot.node(nodeid, label=node['label'], **node['attrs'])
 
-        for node in self.nodes.values():
-            self.dot.node(node.nodeid, label=node.label, **node._attrs)
-
-        for subgraph in self.subgraphs:
-            self.dot.subgraph(subgraph.dot)
+        for dot in self.subgraphs:
+            self.dot.subgraph(dot)
 
         if self._parent:
-            self._parent.remove_node(self.nodeid)
-            self._parent.subgraph(self)
+            self._parent.subgraph(self.dot)
 
-    def node(self, node: "Node") -> None:
+    def node(self, nodeid: str, label: str, **attrs) -> None:
         """Create a new node."""
-        self.nodes[node.nodeid] = node
+        self.nodes[nodeid] = {'label': label, 'attrs': attrs}
 
     def remove_node(self, nodeid: str) -> None:
         del self.nodes[nodeid]
 
-    def subgraph(self, subgraph: "_Cluster") -> None:
+    def subgraph(self, dot: Digraph) -> None:
         """Create a subgraph for clustering"""
-        self.subgraphs.append(subgraph)
-
-    @property
-    def nodes_iter(self):
-        if self.nodes:
-            yield from self.nodes.values()
-        if self.subgraphs:
-            for subgraph in self.subgraphs:
-                yield from subgraph.nodes_iter
+        self.subgraphs.append(dot)
 
     def _validate_direction(self, direction: str):
         direction = direction.upper()
@@ -211,22 +202,23 @@ class Diagram(_Cluster):
     def __enter__(self):
         setdiagram(self)
         super().__enter__()
+        super().__enter__()
         return self
 
-    def __exit__(self, *args):
-        super().__exit__(*args)
+    def __exit__(self, exc_type, exc_value, traceback):
+        super().__exit__(exc_type, exc_value, traceback)
         setdiagram(None)
 
-        for (node1, node2), edge in self.edges.items():
-            cluster_node1 = next(node1.nodes_iter, None)
-            if cluster_node1:
-                edge._attrs['ltail'] = node1.nodeid
-                node1 = cluster_node1
-            cluster_node2 = next(node2.nodes_iter, None)
-            if cluster_node2:
-                edge._attrs['lhead'] = node2.nodeid
-                node2 = cluster_node2
-            self.dot.edge(node1.nodeid, node2.nodeid, **edge.attrs)
+        for nodes, edge in self.edges.items():
+            node1, node2 = nodes
+            nodeid1, nodeid2 = node1.nodeid, node2.nodeid
+            if node1.nodes:
+                edge._attrs['ltail'] = nodeid1
+                nodeid1 = next(iter(node1.nodes.keys()))
+            if node2.nodes:
+                edge._attrs['lhead'] = nodeid2
+                nodeid2 = next(iter(node2.nodes.keys()))
+            self.dot.edge(nodeid1, nodeid2, **edge.attrs)
 
         self.render()
         # Remove the graphviz file leaving only the image.
@@ -378,10 +370,17 @@ class Node(_Cluster):
 
         return self
 
-    def __exit__(self, *args):
-        super().__exit__(*args)
-        self._id = "cluster_" + self.nodeid
-        self.dot.name = self.nodeid
+    def __exit__(self, exc_type, exc_value, traceback):
+        if not (self.nodes or self.subgraphs):
+            return
+
+        self._parent.remove_node(self._id)
+
+        self._id = "cluster_" + self._id
+        self.dot.name = self._id
+
+        super().__exit__(exc_type, exc_value, traceback)
+
 
     def __repr__(self):
         _name = self.__class__.__name__