def search_api_tree(keywords, api_tree):
"""Search Earth Engine API and return functions containing the specified keywords
Args:
keywords (str): The keywords to search for.
api_tree (dict): The dictionary containing the Earth Engine API tree.
Returns:
object: An ipytree object/widget.
"""
import warnings
warnings.filterwarnings("ignore")
sub_tree = Tree()
for key in api_tree.keys():
if keywords.lower() in key.lower():
sub_tree.add_node(api_tree[key])
return sub_tree
def search_api_tree(keywords, api_tree, tools_dict):
"""Search the WhiteboxTools API and return functions containing the specified keywords
Args:
keywords (str): The keywords to search for.
api_tree (dict): The dictionary containing the WhiteboxTools API tree.
tools_dict (dict): The dictionary containing the dict of all tools.
Returns:
object: An ipytree object/widget.
"""
import warnings
warnings.filterwarnings("ignore")
sub_tree = Tree()
for key in api_tree.keys():
if (keywords.lower()
in key.lower()) or (keywords.lower()
in tools_dict[key]["description"].lower()):
sub_tree.add_node(api_tree[key])
return sub_tree
def build_toolbox_tree(tools_dict, folder_icon="folder", tool_icon="wrench"):
"""Build the toolbox for WhiteboxTools.
Args:
tools_dict (dict): A dictionary containing information for all tools.
folder_icon (str, optional): The font-awesome icon for tool categories. Defaults to "folder".
tool_icon (str, optional): The font-awesome icon for tools. Defaults to "wrench".
Returns:
object: An ipywidget representing the toolbox.
"""
left_widget = widgets.VBox()
right_widget = widgets.VBox()
full_widget = widgets.HBox([left_widget, right_widget])
search_description = f"{len(tools_dict)} tools available. Search tools ..."
search_box = widgets.Text(placeholder=search_description)
search_box.layout.width = "270px"
close_btn = widgets.Button(icon="close",
layout=widgets.Layout(width="32px"))
def close_btn_clicked(b):
full_widget.close()
close_btn.on_click(close_btn_clicked)
tree_widget = widgets.Output()
tree_widget.layout.max_width = "310px"
tree_widget.overflow = "auto"
left_widget.children = [widgets.HBox([search_box, close_btn]), tree_widget]
output = widgets.Output(layout=widgets.Layout(max_width="760px"))
right_widget.children = [output]
tree = Tree(multiple_selection=False)
tree_dict = {}
def search_box_callback(text):
with tree_widget:
if text.value == "":
print("Loading...")
tree_widget.clear_output(wait=True)
display(tree)
else:
tree_widget.clear_output()
print("Searching...")
tree_widget.clear_output(wait=True)
sub_tree = search_api_tree(text.value, tree_dict)
display(sub_tree)
search_box.on_submit(search_box_callback)
root_name = "WhiteboxTools"
root_node = Node(root_name)
tree.add_node(root_node)
categories = {}
def handle_tool_clicked(event):
if event["new"]:
cur_node = event["owner"]
tool_name = cur_node.name
with output:
output.clear_output()
tool_ui = tool_gui(tools_dict[tool_name])
display(tool_ui)
for key in tools_dict.keys():
category = tools_dict[key]["category"]
if category not in categories.keys():
category_node = Node(category, icon=folder_icon, opened=False)
root_node.add_node(category_node)
categories[category] = category_node
tool_node = Node(key, icon=tool_icon)
category_node.add_node(tool_node)
tree_dict[key] = tool_node
tool_node.observe(handle_tool_clicked, "selected")
else:
category_node = categories[category]
tool_node = Node(key, icon=tool_icon)
category_node.add_node(tool_node)
tree_dict[key] = tool_node
tool_node.observe(handle_tool_clicked, "selected")
with tree_widget:
tree_widget.clear_output()
display(tree)
return full_widget
def __init__(self, **kwargs):
self._tree = Tree()
self._tree.observe(self._observe_tree_selected_nodes,
["selected_nodes"])
super().__init__(**kwargs)
class NodesTreeWidget(ipw.Output):
"""A tree widget for the structured representation of a nodes graph."""
nodes = traitlets.Tuple().tag(trait=traitlets.Instance(Node))
selected_nodes = traitlets.Tuple(read_only=True).tag(
trait=traitlets.Instance(Node))
PROCESS_STATE_STYLE = {
ProcessState.EXCEPTED: "danger",
ProcessState.FINISHED: "success",
ProcessState.KILLED: "warning",
ProcessState.RUNNING: "info",
ProcessState.WAITING: "info",
}
PROCESS_STATE_STYLE_DEFAULT = "default"
NODE_TYPE = {
WorkChainNode: WorkChainProcessTreeNode,
CalcFunctionNode: CalcFunctionTreeNode,
CalcJobNode: CalcJobTreeNode,
}
def __init__(self, **kwargs):
self._tree = Tree()
self._tree.observe(self._observe_tree_selected_nodes,
["selected_nodes"])
super().__init__(**kwargs)
def _refresh_output(self):
# There appears to be a bug in the ipytree implementation that sometimes
# causes the output to not be properly cleared. We therefore refresh the
# displayed tree upon change of the process trait.
with self:
clear_output()
display(self._tree)
def _observe_tree_selected_nodes(self, change):
return self.set_trait(
"selected_nodes",
tuple(
load_node(pk=node.pk) for node in change["new"]
if hasattr(node, "pk")),
)
def _convert_to_tree_nodes(self, old_nodes, new_nodes):
"Convert nodes into tree nodes while re-using already converted nodes."
old_nodes_ = {node.pk: node for node in old_nodes}
assert len(old_nodes_) == len(old_nodes) # no duplicated nodes
for node in new_nodes:
if node.pk in old_nodes_:
yield old_nodes_[node.pk]
else:
yield self._to_tree_node(node, opened=True)
@traitlets.observe("nodes")
def _observe_nodes(self, change):
self._tree.nodes = list(
sorted(
self._convert_to_tree_nodes(old_nodes=self._tree.nodes,
new_nodes=change["new"]),
key=lambda node: node.pk,
))
self.update()
self._refresh_output()
@classmethod
def _to_tree_node(cls, node, name=None, **kwargs):
"""Convert an AiiDA node to a tree node."""
if name is None:
if isinstance(node, ProcessNode):
name = calc_info(node)
else:
name = str(node)
return cls.NODE_TYPE.get(type(node), UnknownTypeTreeNode)(pk=node.pk,
name=name,
**kwargs)
@classmethod
def _find_called(cls, root):
assert isinstance(root, AiidaProcessNodeTreeNode)
process_node = load_node(root.pk)
called = process_node.called
called.sort(key=lambda p: p.ctime)
for node in called:
if node.pk not in root.nodes_registry:
try:
name = calc_info(node)
except AttributeError:
name = str(node)
root.nodes_registry[node.pk] = cls._to_tree_node(node,
name=name)
yield root.nodes_registry[node.pk]
@classmethod
def _find_outputs(cls, root):
process_node = load_node(root.parent_pk)
if root.namespace:
outputs = process_node.outputs[root.namespace]
else:
outputs = process_node.outputs
output_nodes = {key: outputs[key] for key in outputs}
for key in sorted(output_nodes.keys(),
key=lambda k: getattr(outputs[k], "pk", -1)):
node = output_nodes[key]
if isinstance(node, AttributeDict):
yield AiidaOutputsTreeNode(name=key,
parent_pk=root.parent_pk,
namespace=key)
else:
if node.pk not in root.nodes_registry:
root.nodes_registry[node.pk] = cls._to_tree_node(
node, name=f"{key}<{node.pk}>")
yield root.nodes_registry[node.pk]
@classmethod
def _find_children(cls, root):
"""Find all children of the provided AiiDA node."""
if isinstance(root, AiidaProcessNodeTreeNode):
yield root.outputs_node
yield from cls._find_called(root)
elif isinstance(root, AiidaOutputsTreeNode):
yield from cls._find_outputs(root)
@classmethod
def _build_tree(cls, root):
"""Recursively build a tree nodes graph for a given tree node."""
root.nodes = [
cls._build_tree(child) for child in cls._find_children(root)
]
return root
@classmethod
def _walk_tree(cls, root):
"""Breadth-first search of the node tree."""
yield root
for node in root.nodes:
yield from cls._walk_tree(node)
def _update_tree_node(self, tree_node):
if isinstance(tree_node, AiidaProcessNodeTreeNode):
process_node = load_node(tree_node.pk)
tree_node.name = calc_info(process_node)
tree_node.icon_style = self.PROCESS_STATE_STYLE.get(
process_node.process_state, self.PROCESS_STATE_STYLE_DEFAULT)
def update(self, _=None):
"""Refresh nodes based on the latest state of the root process and its children."""
for root_node in self._tree.nodes:
self._build_tree(root_node)
for tree_node in self._walk_tree(root_node):
self._update_tree_node(tree_node)
def find_node(self, pk):
for node in self._walk_tree(self._tree):
if getattr(node, "pk", None) == pk:
return node
raise KeyError(pk)
def test_tree():
Tree()