diff --git a/src/ctm_python_client/ext/viz.py b/src/ctm_python_client/ext/viz.py
index 9987b7f..633f9f9 100644
--- a/src/ctm_python_client/ext/viz.py
+++ b/src/ctm_python_client/ext/viz.py
@@ -1,55 +1,66 @@
 import graphviz
-from aapi import Folder, SubFolder, SimpleFolder
+
+from aapi import Folder, SimpleFolder, SubFolder
 from ctm_python_client.core.workflow import BaseWorkflow
 
 
 def get_subgraph(obj, name, current_path):
-    if not(
-        isinstance(obj, Folder) or
-        isinstance(obj, SubFolder) or
-        isinstance(obj, SimpleFolder)
+    if not (
+        isinstance(obj, Folder)
+        or isinstance(obj, SubFolder)
+        or isinstance(obj, SimpleFolder)
     ):
         return
-    sgraph = graphviz.Digraph(name, graph_attr={'label':obj.object_name}, node_attr={'shape':'oval'})
+    sgraph = graphviz.Digraph(
+        name, graph_attr={"label": obj.object_name}, node_attr={"shape": "oval"}
+    )
     for o in obj.job_list:
-        sgraph.node(f'{current_path}/{o.object_name}', o.object_name)
+        sgraph.node(f"{current_path}/{o.object_name}", o.object_name)
 
-    for i,o in enumerate(obj.sub_folder_list):
-        sgraph.subgraph(get_subgraph(o, f'cluster_{current_path}_{i}', current_path+'/'+o.object_name))
+    for i, o in enumerate(obj.sub_folder_list):
+        sgraph.subgraph(
+            get_subgraph(
+                o, f"cluster_{current_path}_{i}", current_path + "/" + o.object_name
+            )
+        )
 
     return sgraph
 
-def get_graph(workflow : BaseWorkflow):
-    dgraph = graphviz.Digraph(name='root', graph_attr={
-        'rankdir': 'LR',
-        'compound': 'true'
-    },
-    node_attr={'shape':'oval'}
+
+def get_graph(workflow: BaseWorkflow):
+    dgraph = graphviz.Digraph(
+        name="root",
+        graph_attr={"rankdir": "LR", "compound": "true"},
+        node_attr={"shape": "oval"},
     )
     i = 0
     for k, v in workflow._definitions.items():
-        if not(
-            isinstance(v, Folder) or
-            isinstance(v, SubFolder) or
-            isinstance(v, SimpleFolder)
+        if not (
+            isinstance(v, Folder)
+            or isinstance(v, SubFolder)
+            or isinstance(v, SimpleFolder)
         ):
             continue
-        cluster = graphviz.Digraph(f'cluster_{i}', graph_attr={
-            'label': v.object_name
-        })
+        cluster = graphviz.Digraph(f"cluster_{i}", graph_attr={"label": v.object_name})
 
         for o in v.job_list:
-            cluster.node(f'{v.object_name}/{o.object_name}', o.object_name)
+            cluster.node(f"{v.object_name}/{o.object_name}", o.object_name)
 
-        for i,o in enumerate(v.sub_folder_list):
-            cluster.subgraph(get_subgraph(o,f'cluster_{v.object_name}/{o.object_name}_{i}',v.object_name+'/'+o.object_name))
-        
+        if hasattr(v, "sub_folder_list"):
+            for i, o in enumerate(v.sub_folder_list):
+                cluster.subgraph(
+                    get_subgraph(
+                        o,
+                        f"cluster_{v.object_name}/{o.object_name}_{i}",
+                        v.object_name + "/" + o.object_name,
+                    )
+                )
 
         dgraph.subgraph(cluster)
         i += 1
 
     for conn in workflow._connections.values():
         for o in conn:
-            dgraph.edge(o[0],o[1])
+            dgraph.edge(o[0], o[1])
 
     return dgraph