@@ -372,7 +372,20 @@ def edges(self):
372372 return SequenceEdge .from_edges ("eval" , out_edges , output )
373373
374374 def visualize (self , input : dict [str , Desc ] | None = None ):
375- import networkx as nx
375+ if input is None :
376+ from .introspection import draw_graph
377+
378+ draw_graph (self )
379+ return
380+
381+ try :
382+ import networkx as nx
383+ except ImportError :
384+ from .introspection import draw_graph
385+
386+ draw_graph (self )
387+ return
388+
376389 import matplotlib .pyplot as plt
377390
378391 from pprint import pformat
@@ -382,38 +395,26 @@ def node_format(x):
382395
383396 G = nx .DiGraph ()
384397
385- if input is not None :
386- for _ , edges in self ._subgraphs :
387- q : list [dict [str , Desc ]] = [input ]
388- explored : set [tuple [tuple [str , str ], ...]] = set ()
389- explored .add (
390- tuple (sorted (((k , v .coordinates ) for k , v in q [0 ].items ())))
391- )
392- G .add_node (node_format (q [0 ]))
393- while q :
394- n = q .pop ()
395- for e in edges :
396- if Desc .compatible (n , e .input ):
397- w = n | e .output
398- if node_format (w ) not in G :
399- G .add_node (node_format (w ))
400- explored .add (
401- tuple (
402- sorted (
403- ((k , v .coordinates ) for k , v in w .items ())
404- )
405- )
398+ for _ , edges in self ._subgraphs :
399+ q : list [dict [str , Desc ]] = [input ]
400+ explored : set [tuple [tuple [str , str ], ...]] = set ()
401+ explored .add (tuple (sorted (((k , v .coordinates ) for k , v in q [0 ].items ()))))
402+ G .add_node (node_format (q [0 ]))
403+ while q :
404+ n = q .pop ()
405+ for e in edges :
406+ if Desc .compatible (n , e .input ):
407+ w = n | e .output
408+ if node_format (w ) not in G :
409+ G .add_node (node_format (w ))
410+ explored .add (
411+ tuple (
412+ sorted (((k , v .coordinates ) for k , v in w .items ()))
406413 )
407- q .append (w )
408- if node_format (w ) != node_format (n ):
409- G .add_edge (node_format (n ), node_format (w ), name = e .name )
410- else :
411- # don't bother separating subgraphs,as the end result is exactly the same here
412- for edge in self ._edges :
413- G .add_edge (
414- node_format (edge .input ), node_format (edge .output ), name = edge .name
415- )
416-
414+ )
415+ q .append (w )
416+ if node_format (w ) != node_format (n ):
417+ G .add_edge (node_format (n ), node_format (w ), name = e .name )
417418 try :
418419 pos = nx .shell_layout (G )
419420 except Exception :
0 commit comments