55import os
66import re
77
8- from Qt import QtCore , QtWidgets , QtGui
8+ from Qt import QtCore , QtWidgets
99
1010from NodeGraphQt .base .commands import (NodeAddedCmd ,
1111 NodeRemovedCmd ,
1818from NodeGraphQt .base .port import Port
1919from NodeGraphQt .constants import (
2020 NODE_LAYOUT_DIRECTION , NODE_LAYOUT_HORIZONTAL , NODE_LAYOUT_VERTICAL ,
21- PIPE_LAYOUT_CURVED , PIPE_LAYOUT_STRAIGHT , PIPE_LAYOUT_ANGLE ,
21+ PipeLayoutEnum ,
2222 URI_SCHEME , URN_SCHEME ,
23- IN_PORT , OUT_PORT ,
24- VIEWER_GRID_LINES
23+ PortTypeEnum ,
24+ ViewerEnum
2525)
2626from NodeGraphQt .nodes .backdrop_node import BackdropNode
2727from NodeGraphQt .nodes .base_node import BaseNode
@@ -369,7 +369,8 @@ def _on_connection_changed(self, disconnected, connected):
369369 return
370370
371371 label = 'connect node(s)' if connected else 'disconnect node(s)'
372- ptypes = {IN_PORT : 'inputs' , OUT_PORT : 'outputs' }
372+ ptypes = {PortTypeEnum .IN .value : 'inputs' ,
373+ PortTypeEnum .OUT .value : 'outputs' }
373374
374375 self ._undo_stack .beginMacro (label )
375376 for p1_view , p2_view in disconnected :
@@ -396,7 +397,8 @@ def _on_connection_sliced(self, ports):
396397 """
397398 if not ports :
398399 return
399- ptypes = {IN_PORT : 'inputs' , OUT_PORT : 'outputs' }
400+ ptypes = {PortTypeEnum .IN .value : 'inputs' ,
401+ PortTypeEnum .OUT .value : 'outputs' }
400402 self ._undo_stack .beginMacro ('slice connections' )
401403 for p1_view , p2_view in ports :
402404 node1 = self ._model .nodes [p1_view .node .id ]
@@ -553,7 +555,7 @@ def set_grid_color(self, r, g, b):
553555 self .scene ().grid_color = (r , g , b )
554556 self ._viewer .force_update ()
555557
556- def set_grid_mode (self , mode = VIEWER_GRID_LINES ):
558+ def set_grid_mode (self , mode = None ):
557559 """
558560 Set node graph grid mode.
559561
@@ -562,13 +564,20 @@ def set_grid_mode(self, mode=VIEWER_GRID_LINES):
562564
563565 Node graph background types:
564566
565- * :attr:`NodeGraphQt.constants.VIEWER_GRID_NONE `
566- * :attr:`NodeGraphQt.constants.VIEWER_GRID_DOTS `
567- * :attr:`NodeGraphQt.constants.VIEWER_GRID_LINES `
567+ * :attr:`NodeGraphQt.constants.ViewerEnum.GRID_DISPLAY_NONE.value `
568+ * :attr:`NodeGraphQt.constants.ViewerEnum.GRID_DISPLAY_DOTS.value `
569+ * :attr:`NodeGraphQt.constants.ViewerEnum.GRID_DISPLAY_LINES.value `
568570
569571 Args:
570572 mode (int): background style.
571573 """
574+ display_types = [
575+ ViewerEnum .GRID_DISPLAY_NONE .value ,
576+ ViewerEnum .GRID_DISPLAY_DOTS .value ,
577+ ViewerEnum .GRID_DISPLAY_LINES .value
578+ ]
579+ if mode not in display_types :
580+ mode = ViewerEnum .GRID_DISPLAY_LINES .value
572581 self .scene ().grid_mode = mode
573582 self ._viewer .force_update ()
574583
@@ -752,7 +761,7 @@ def set_pipe_collision(self, mode=True):
752761 self ._model .pipe_collision = mode
753762 self ._viewer .pipe_collision = mode
754763
755- def set_pipe_style (self , style = PIPE_LAYOUT_CURVED ):
764+ def set_pipe_style (self , style = PipeLayoutEnum . CURVED . value ):
756765 """
757766 Set node graph pipes to be drawn as straight, curved or angled.
758767
@@ -764,17 +773,17 @@ def set_pipe_style(self, style=PIPE_LAYOUT_CURVED):
764773
765774 Pipe Layout Styles:
766775
767- * :attr:`NodeGraphQt.constants.PIPE_LAYOUT_CURVED `
768- * :attr:`NodeGraphQt.constants.PIPE_LAYOUT_STRAIGHT `
769- * :attr:`NodeGraphQt.constants.PIPE_LAYOUT_ANGLE `
776+ * :attr:`NodeGraphQt.constants.PipeLayoutEnum.CURVED.value `
777+ * :attr:`NodeGraphQt.constants.PipeLayoutEnum.STRAIGHT.value `
778+ * :attr:`NodeGraphQt.constants.PipeLayoutEnum.ANGLE.value `
770779
771780 Args:
772781 style (int): pipe layout style.
773782 """
774- pipe_max = max ([PIPE_LAYOUT_CURVED ,
775- PIPE_LAYOUT_STRAIGHT ,
776- PIPE_LAYOUT_ANGLE ])
777- style = style if 0 <= style <= pipe_max else PIPE_LAYOUT_CURVED
783+ pipe_max = max ([PipeLayoutEnum . CURVED . value ,
784+ PipeLayoutEnum . STRAIGHT . value ,
785+ PipeLayoutEnum . ANGLE . value ])
786+ style = style if 0 <= style <= pipe_max else PipeLayoutEnum . CURVED . value
778787 self ._viewer .set_pipe_layout (style )
779788
780789 def fit_to_selection (self ):
@@ -1246,16 +1255,20 @@ def _serialize(self, nodes):
12461255 for pname , conn_data in inputs .items ():
12471256 for conn_id , prt_names in conn_data .items ():
12481257 for conn_prt in prt_names :
1249- pipe = {IN_PORT : [n_id , pname ],
1250- OUT_PORT : [conn_id , conn_prt ]}
1258+ pipe = {
1259+ PortTypeEnum .IN .value : [n_id , pname ],
1260+ PortTypeEnum .OUT .value : [conn_id , conn_prt ]
1261+ }
12511262 if pipe not in serial_data ['connections' ]:
12521263 serial_data ['connections' ].append (pipe )
12531264
12541265 for pname , conn_data in outputs .items ():
12551266 for conn_id , prt_names in conn_data .items ():
12561267 for conn_prt in prt_names :
1257- pipe = {OUT_PORT : [n_id , pname ],
1258- IN_PORT : [conn_id , conn_prt ]}
1268+ pipe = {
1269+ PortTypeEnum .OUT .value : [n_id , pname ],
1270+ PortTypeEnum .IN .value : [conn_id , conn_prt ]
1271+ }
12591272 if pipe not in serial_data ['connections' ]:
12601273 serial_data ['connections' ].append (pipe )
12611274
@@ -2362,8 +2375,10 @@ def get_node_by_port(self, port):
23622375 Returns:
23632376 PortInputNode or PortOutputNode: port node object.
23642377 """
2365- func_type = {IN_PORT : self .get_input_port_nodes ,
2366- OUT_PORT : self .get_output_port_nodes }
2378+ func_type = {
2379+ PortTypeEnum .IN .value : self .get_input_port_nodes ,
2380+ PortTypeEnum .OUT .value : self .get_output_port_nodes
2381+ }
23672382 for n in func_type .get (port .type_ (), []):
23682383 if port == n .parent_port :
23692384 return n
0 commit comments