@@ -361,7 +361,20 @@ def edges(self):
361361 return SequenceEdge .from_edges ("eval" , out_edges , output )
362362
363363 def visualize (self , input : dict [str , Desc ] | None = None ):
364- import networkx as nx
364+ if input is None :
365+ from .introspection import draw_graph
366+
367+ draw_graph (self )
368+ return
369+
370+ try :
371+ import networkx as nx
372+ except ImportError :
373+ from .introspection import draw_graph
374+
375+ draw_graph (self )
376+ return
377+
365378 import matplotlib .pyplot as plt
366379
367380 from pprint import pformat
@@ -371,39 +384,26 @@ def node_format(x):
371384
372385 G = nx .DiGraph ()
373386
374- if input is not None :
375-
376- for _ , edges in self ._subgraphs :
377- q : list [dict [str , Desc ]] = [input ]
378- explored : set [tuple [tuple [str , str ], ...]] = set ()
379- explored .add (
380- tuple (sorted (((k , v .coordinates ) for k , v in q [0 ].items ())))
381- )
382- G .add_node (node_format (q [0 ]))
383- while q :
384- n = q .pop ()
385- for e in edges :
386- if Desc .compatible (n , e .input ):
387- w = n | e .output
388- if node_format (w ) not in G :
389- G .add_node (node_format (w ))
390- explored .add (
391- tuple (
392- sorted (
393- ((k , v .coordinates ) for k , v in w .items ())
394- )
395- )
387+ for _ , edges in self ._subgraphs :
388+ q : list [dict [str , Desc ]] = [input ]
389+ explored : set [tuple [tuple [str , str ], ...]] = set ()
390+ explored .add (tuple (sorted (((k , v .coordinates ) for k , v in q [0 ].items ()))))
391+ G .add_node (node_format (q [0 ]))
392+ while q :
393+ n = q .pop ()
394+ for e in edges :
395+ if Desc .compatible (n , e .input ):
396+ w = n | e .output
397+ if node_format (w ) not in G :
398+ G .add_node (node_format (w ))
399+ explored .add (
400+ tuple (
401+ sorted (((k , v .coordinates ) for k , v in w .items ()))
396402 )
397- q .append (w )
398- if node_format (w ) != node_format (n ):
399- G .add_edge (node_format (n ), node_format (w ), name = e .name )
400- else :
401- # don't bother separating subgraphs,as the end result is exactly the same here
402- for edge in self ._edges :
403- G .add_edge (
404- node_format (edge .input ), node_format (edge .output ), name = edge .name
405- )
406-
403+ )
404+ q .append (w )
405+ if node_format (w ) != node_format (n ):
406+ G .add_edge (node_format (n ), node_format (w ), name = e .name )
407407 try :
408408 pos = nx .shell_layout (G )
409409 except Exception :
0 commit comments