diff --git a/NodeGraphQt/base/graph.py b/NodeGraphQt/base/graph.py index 887bd971..f7e67c22 100644 --- a/NodeGraphQt/base/graph.py +++ b/NodeGraphQt/base/graph.py @@ -17,12 +17,13 @@ from .model import NodeGraphModel from .node import NodeObject, BaseNode, BackdropNode from .port import Port -from ..constants import (URI_SCHEME, URN_SCHEME, - PIPE_LAYOUT_CURVED, - PIPE_LAYOUT_STRAIGHT, - PIPE_LAYOUT_ANGLE, - IN_PORT, OUT_PORT, - VIEWER_GRID_LINES) +from ..constants import ( + URI_SCHEME, URN_SCHEME, + NODE_LAYOUT_DIRECTION, NODE_LAYOUT_HORIZONTAL, NODE_LAYOUT_VERTICAL, + PIPE_LAYOUT_CURVED, PIPE_LAYOUT_STRAIGHT, PIPE_LAYOUT_ANGLE, + IN_PORT, OUT_PORT, + VIEWER_GRID_LINES +) from ..widgets.node_space_bar import node_space_bar from ..widgets.viewer import NodeViewer @@ -1558,6 +1559,149 @@ def disable_nodes(self, nodes, mode=None): return nodes[0].set_disabled(mode) + # auto layout node functions. + + @staticmethod + def _update_node_rank(node, nodes_rank, down_stream=True): + """ + Recursive function for updating the node ranking. + + Args: + node (NodeGraphQt.BaseNode): node to start from. + nodes_rank (dict): node ranking object to be updated. + down_stream (bool): true to rank down stram. + """ + if down_stream: + node_values = node.connected_output_nodes().values() + else: + node_values = node.connected_input_nodes().values() + + connected_nodes = set() + for nodes in node_values: + connected_nodes.update(nodes) + + rank = nodes_rank[node] + 1 + for n in connected_nodes: + if n in nodes_rank: + nodes_rank[n] = max(nodes_rank[n], rank) + else: + nodes_rank[n] = rank + NodeGraph._update_node_rank(n, nodes_rank, down_stream) + + @staticmethod + def _compute_node_rank(nodes, down_stream=True): + """ + Compute the ranking of nodes. + + Args: + nodes (list[NodeGraphQt.BaseNode]): nodes to start ranking from. + down_stream (bool): true to compute down stream. + + Returns: + dict: {NodeGraphQt.BaseNode: node_rank, ...} + """ + nodes_rank = {} + for node in nodes: + nodes_rank[node] = 0 + NodeGraph._update_node_rank(node, nodes_rank, down_stream) + return nodes_rank + + def auto_layout_nodes(self, nodes=None, down_stream=True, start_nodes=None): + """ + Auto layout the nodes in the node graph. + + Note: + If the node graph is acyclic then the ``start_nodes`` will need + to be specified. + + Args: + nodes (list[NodeGraphQt.BaseNode]): list of nodes to auto layout + if nodes is None then all nodes is layed out. + down_stream (bool): false to layout up stream. + start_nodes (list[NodeGraphQt.BaseNode]): + list of nodes to start the auto layout from (Optional). + """ + self.begin_undo('Auto Layout Nodes') + + nodes = nodes or self.all_nodes() + + # filter out the backdrops. + backdrops = { + n: n.nodes() for n in nodes if isinstance(n, BackdropNode) + } + filtered_nodes = [n for n in nodes if not isinstance(n, BackdropNode)] + + start_nodes = start_nodes or [] + if down_stream: + start_nodes += [ + n for n in filtered_nodes + if not any(n.connected_input_nodes().values()) + ] + else: + start_nodes += [ + n for n in filtered_nodes + if not any(n.connected_output_nodes().values()) + ] + + if not start_nodes: + return + + node_views = [n.view for n in nodes] + nodes_center_0 = self.viewer().nodes_rect_center(node_views) + + nodes_rank = NodeGraph._compute_node_rank(start_nodes, down_stream) + + rank_map = {} + for node, rank in nodes_rank.items(): + if rank in rank_map: + rank_map[rank].append(node) + else: + rank_map[rank] = [node] + + if NODE_LAYOUT_DIRECTION is NODE_LAYOUT_HORIZONTAL: + current_x = 0 + node_height = 120 + for rank in sorted(range(len(rank_map)), reverse=not down_stream): + ranked_nodes = rank_map[rank] + max_width = max([node.view.width for node in ranked_nodes]) + current_x += max_width + current_y = 0 + for idx, node in enumerate(ranked_nodes): + dy = max(node_height, node.view.height) + current_y += 0 if idx == 0 else dy + node.set_pos(current_x, current_y) + current_y += dy * 0.5 + 10 + + current_x += max_width * 0.5 + 100 + elif NODE_LAYOUT_DIRECTION is NODE_LAYOUT_VERTICAL: + current_y = 0 + node_width = 250 + for rank in sorted(range(len(rank_map)), reverse=not down_stream): + ranked_nodes = rank_map[rank] + max_height = max([node.view.height for node in ranked_nodes]) + current_y += max_height + current_x = 0 + for idx, node in enumerate(ranked_nodes): + dx = max(node_width, node.view.width) + current_x += 0 if idx == 0 else dx + node.set_pos(current_x, current_y) + current_x += dx * 0.5 + 10 + + current_y += max_height * 0.5 + 100 + + nodes_center_1 = self.viewer().nodes_rect_center(node_views) + dx = nodes_center_0[0] - nodes_center_1[0] + dy = nodes_center_0[1] - nodes_center_1[1] + [n.set_pos(n.x_pos() + dx, n.y_pos() + dy) for n in nodes] + + # wrap the backdrop nodes. + for backdrop, contained_nodes in backdrops.items(): + backdrop.wrap_nodes(contained_nodes) + + self.end_undo() + + # prompt dialog functions. + def question_dialog(self, text, title='Node Graph'): """ Prompts a question open dialog with ``"Yes"`` and ``"No"`` buttons in @@ -1624,6 +1768,8 @@ def save_dialog(self, current_dir=None, ext=None): """ return self._viewer.save_dialog(current_dir, ext) + ### --- + def use_OpenGL(self): """ Use OpenGL to draw the graph. diff --git a/NodeGraphQt/base/utils.py b/NodeGraphQt/base/utils.py index 000381c3..78187336 100644 --- a/NodeGraphQt/base/utils.py +++ b/NodeGraphQt/base/utils.py @@ -87,9 +87,9 @@ def setup_context_menu(graph): edit_menu.add_separator() edit_menu.add_command( - 'Layout Nodes Up Stream', _layout_graph_up, 'L') + 'Auto Layout Up Stream', _layout_graph_up, 'L') edit_menu.add_command( - 'Layout Nodes Down Stream', _layout_graph_down, 'Ctrl+L') + 'Auto Layout Down Stream', _layout_graph_down, 'Ctrl+L') edit_menu.add_separator() @@ -302,34 +302,15 @@ def _bg_grid_lines(graph): graph.set_grid_mode(2) -def __layout_graph(graph, down_stream=True): - graph.begin_undo('Auto Layout') - node_space = graph.get_node_space() - if node_space is not None: - nodes = node_space.children() - else: - nodes = graph.all_nodes() - if not nodes: - return - node_views = [n.view for n in nodes] - nodes_center0 = graph.viewer().nodes_rect_center(node_views) - if down_stream: - auto_layout_down(all_nodes=nodes) - else: - auto_layout_up(all_nodes=nodes) - nodes_center1 = graph.viewer().nodes_rect_center(node_views) - dx = nodes_center0[0] - nodes_center1[0] - dy = nodes_center0[1] - nodes_center1[1] - [n.set_pos(n.x_pos() + dx, n.y_pos()+dy) for n in nodes] - graph.end_undo() - - def _layout_graph_down(graph): - __layout_graph(graph, True) + nodes = graph.selected_nodes() or graph.all_nodes() + graph.auto_layout_nodes(nodes=nodes, down_stream=True) def _layout_graph_up(graph): - __layout_graph(graph, False) + nodes = graph.selected_nodes() or graph.all_nodes() + graph.auto_layout_nodes(nodes-nodes, down_stream=False) + # topological_sort @@ -628,194 +609,6 @@ def update_nodes_by_up(nodes): _update_nodes(topological_sort_by_up(all_nodes=nodes)) -# auto layout - - -def _update_node_rank_down(node, nodes_rank): - rank = nodes_rank[node] + 1 - for n in get_output_nodes(node, False): - if n in nodes_rank: - nodes_rank[n] = max(nodes_rank[n], rank) - else: - nodes_rank[n] = rank - _update_node_rank_down(n, nodes_rank) - - -def _compute_rank_down(start_nodes): - """ - Compute the rank of the down stream nodes. - - Args: - start_nodes (list[NodeGraphQt.BaseNode]): - (Optional) the start nodes of the graph. - - Returns: - dict{NodeGraphQt.BaseNode: node_rank, ...} - """ - nodes_rank = {} - for node in start_nodes: - nodes_rank[node] = 0 - _update_node_rank_down(node, nodes_rank) - return nodes_rank - - -def _update_node_rank_up(node, nodes_rank): - rank = nodes_rank[node] + 1 - for n in get_input_nodes(node): - if n in nodes_rank: - nodes_rank[n] = max(nodes_rank[n], rank) - else: - nodes_rank[n] = rank - _update_node_rank_up(n, nodes_rank) - - -def _compute_rank_up(start_nodes): - """ - Compute the rank of the up stream nodes. - - Args: - start_nodes (list[NodeGraphQt.BaseNode]): - (Optional) the end nodes of the graph. - - Returns: - dict{NodeGraphQt.BaseNode: node_rank, ...} - """ - - nodes_rank = {} - for node in start_nodes: - nodes_rank[node] = 0 - _update_node_rank_up(node, nodes_rank) - return nodes_rank - - -def auto_layout_up(start_nodes=None, all_nodes=None): - """ - Auto layout the nodes by up stream direction. - - Args: - start_nodes (list[NodeGraphQt.BaseNode]): - (Optional) the end nodes of the graph. - all_nodes (list[NodeGraphQt.BaseNode]): - (Optional) if 'start_nodes' is None the function can calculate - start nodes from 'all_nodes'. - """ - if not start_nodes and not all_nodes: - return - if start_nodes: - start_nodes = __remove_BackdropNode(start_nodes) - if all_nodes: - all_nodes = __remove_BackdropNode(all_nodes) - - if not start_nodes: - start_nodes = [n for n in all_nodes if not _has_output_node(n)] - if not start_nodes: - return - - nodes_rank = _compute_rank_up(start_nodes) - - rank_map = {} - for node, rank in nodes_rank.items(): - if rank in rank_map: - rank_map[rank].append(node) - else: - rank_map[rank] = [node] - - if NODE_LAYOUT_DIRECTION is NODE_LAYOUT_HORIZONTAL: - current_x = 0 - node_height = 80 - for rank in reversed(range(len(rank_map))): - nodes = rank_map[rank] - max_width = max([node.view.width for node in nodes]) - current_x += max_width - current_y = 0 - for idx, node in enumerate(nodes): - dy = max(node_height, node.view.height) - current_y += 0 if idx == 0 else dy - node.set_pos(current_x, current_y) - current_y += dy * 0.5 + 10 - - current_x += max_width * 0.5 + 100 - elif NODE_LAYOUT_DIRECTION is NODE_LAYOUT_VERTICAL: - current_y = 0 - node_width = 250 - for rank in reversed(range(len(rank_map))): - nodes = rank_map[rank] - max_height = max([node.view.height for node in nodes]) - current_y += max_height - current_x = 0 - for idx, node in enumerate(nodes): - dx = max(node_width, node.view.width) - current_x += 0 if idx == 0 else dx - node.set_pos(current_x, current_y) - current_x += dx * 0.5 + 10 - - current_y += max_height * 0.5 + 100 - - -def auto_layout_down(start_nodes=None, all_nodes=None): - """ - Auto layout the nodes by down stream direction. - - Args: - start_nodes (list[NodeGraphQt.BaseNode]): - (Optional) the start update nodes of the graph. - all_nodes (list[NodeGraphQt.BaseNode]): - (Optional) if 'start_nodes' is None the function can calculate - start nodes from 'all_nodes'. - """ - if not start_nodes and not all_nodes: - return - if start_nodes: - start_nodes = __remove_BackdropNode(start_nodes) - if all_nodes: - all_nodes = __remove_BackdropNode(all_nodes) - - if not start_nodes: - start_nodes = [n for n in all_nodes if not _has_input_node(n)] - if not start_nodes: - return - - nodes_rank = _compute_rank_down(start_nodes) - - rank_map = {} - for node, rank in nodes_rank.items(): - if rank in rank_map: - rank_map[rank].append(node) - else: - rank_map[rank] = [node] - - if NODE_LAYOUT_DIRECTION is NODE_LAYOUT_HORIZONTAL: - current_x = 0 - node_height = 120 - for rank in range(len(rank_map)): - nodes = rank_map[rank] - max_width = max([node.view.width for node in nodes]) - current_x += max_width - current_y = 0 - for idx, node in enumerate(nodes): - dy = max(node_height, node.view.height) - current_y += 0 if idx == 0 else dy - node.set_pos(current_x, current_y) - current_y += dy * 0.5 + 10 - - current_x += max_width * 0.5 + 100 - elif NODE_LAYOUT_DIRECTION is NODE_LAYOUT_VERTICAL: - current_y = 0 - node_width = 250 - for rank in range(len(rank_map)): - nodes = rank_map[rank] - max_height = max([node.view.height for node in nodes]) - current_y += max_height - current_x = 0 - for idx, node in enumerate(nodes): - dx = max(node_width, node.view.width) - current_x += 0 if idx == 0 else dx - node.set_pos(current_x, current_y) - current_x += dx * 0.5 + 10 - - current_y += max_height * 0.5 + 100 - - # garbage collect def minimize_node_ref_count(node): diff --git a/README.md b/README.md index 624d515e..9a6e651a 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,10 @@ if __name__ == '__main__': node_b = graph.create_node('com.chantasticvfx.MyNode', name='Node B', color='#5b162f') # connect node a input to node b output. - node_a.set_input(0, node_b.output(0)) + node_a.set_input(0, node_b.output(0)) + + # auto layout nodes. + graph.auto_layout_nodes() # get the widget and show. graph_widget = graph.widget diff --git a/example.py b/example.py index f002ae79..c6ef5f86 100644 --- a/example.py +++ b/example.py @@ -178,31 +178,32 @@ def show_nodes_list(node): ] graph.register_nodes(nodes_to_reg) - my_node = graph.create_node('com.chantasticvfx.MyNode', - name='chantastic!', - color='#0a1e20', - text_color='#feab20', - pos=[310, 10]) - - foo_node = graph.create_node('com.chantasticvfx.FooNode', - name='node', - pos=[-480, 140]) + my_node = graph.create_node( + 'com.chantasticvfx.MyNode', + name='chantastic!', + color='#0a1e20', + text_color='#feab20' + ) + + foo_node = graph.create_node( + 'com.chantasticvfx.FooNode', + name='node') foo_node.set_disabled(True) # create example "TextInputNode". - text_node = graph.create_node('com.chantasticvfx.TextInputNode', - name='text node', - pos=[-480, -160]) + text_node = graph.create_node( + 'com.chantasticvfx.TextInputNode', + name='text node') # create example "TextInputNode". - checkbox_node = graph.create_node('com.chantasticvfx.CheckboxNode', - name='checkbox node', - pos=[-480, -20]) + checkbox_node = graph.create_node( + 'com.chantasticvfx.CheckboxNode', + name='checkbox node') # create node with a combo box menu. - menu_node = graph.create_node('com.chantasticvfx.DropdownMenuNode', - name='menu node', - pos=[280, -200]) + menu_node = graph.create_node( + 'com.chantasticvfx.DropdownMenuNode', + name='menu node') # change node icon. this_path = os.path.dirname(os.path.abspath(__file__)) @@ -210,16 +211,18 @@ def show_nodes_list(node): bar_node = graph.create_node('com.chantasticvfx.BarNode') bar_node.set_icon(icon) bar_node.set_name('icon node') - bar_node.set_pos(-70, 10) # connect the nodes. foo_node.set_output(0, bar_node.input(2)) menu_node.set_input(0, bar_node.output(1)) bar_node.set_input(0, text_node.output(0)) + # auto layout nodes. + graph.auto_layout_nodes() + # wrap a backdrop node. backdrop_node = graph.create_node('nodeGraphQt.nodes.BackdropNode') - backdrop_node.wrap_nodes([my_node, menu_node]) + backdrop_node.wrap_nodes([text_node, checkbox_node]) graph.fit_to_selection()