# -*- coding: utf-8 -*- ########################################################################### # Copyright (c), The AiiDA team. All rights reserved. # # This file is part of the AiiDA code. # # # # The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### """Utility functions to draw ASCII diagrams to the command line.""" from aiida.common.links import LinkType __all__ = ('draw_children', 'draw_parents', 'format_call_graph') TREE_LAST_ENTRY = '\u2514\u2500\u2500 ' TREE_MIDDLE_ENTRY = '\u251C\u2500\u2500 ' TREE_FIRST_ENTRY = TREE_MIDDLE_ENTRY class NodeTreePrinter: """Utility functions for printing node trees. .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. """ # Note: when removing this code, also remove the `ete3` as a dependency as it will no longer be used. @classmethod def print_node_tree(cls, node, max_depth, follow_links=()): """Top-level function for printing node tree.""" import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('class is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member from ete3 import Tree from aiida.cmdline.utils.common import get_node_summary from aiida.cmdline.utils import echo echo.echo(get_node_summary(node)) tree_string = f'({cls._build_tree(node, max_depth=max_depth, follow_links=follow_links)});' tmp = Tree(tree_string, format=1) echo.echo(tmp.get_ascii(show_internal=True)) @staticmethod def _ctime(link_triple): return link_triple.node.ctime @classmethod def _build_tree(cls, node, show_pk=True, max_depth=None, follow_links=(), depth=0): """Return string with tree.""" if max_depth is not None and depth > max_depth: return None children = [] for entry in sorted(node.get_outgoing(link_type=follow_links).all(), key=cls._ctime): child_str = cls._build_tree( entry.node, show_pk, follow_links=follow_links, max_depth=max_depth, depth=depth + 1 ) if child_str: children.append(child_str) out_values = [] if children: out_values.append('(') out_values.append(', '.join(children)) out_values.append(')') lab = node.__class__.__name__ if show_pk: lab += f' [{node.pk}]' out_values.append(lab) return ''.join(out_values) [docs]def draw_parents(node, node_label=None, show_pk=True, dist=2, follow_links_of_type=None): """ Print an ASCII tree of the parents of the given node. .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. :param node: The node to draw for :type node: :class:`aiida.orm.nodes.data.Data` :param node_label: The label to use for the nodes :type node_label: str :param show_pk: Show the PK of nodes alongside the label :type show_pk: bool :param dist: The number of steps away from this node to branch out :type dist: int :param follow_links_of_type: Follow links of this type when making steps, if None then it will follow CREATE and INPUT links :type follow_links_of_type: str """ import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('function is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member return get_ascii_tree(node, node_label, show_pk, dist, follow_links_of_type, False) [docs]def draw_children(node, node_label=None, show_pk=True, dist=2, follow_links_of_type=None): """ Print an ASCII tree of the parents of the given node. .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. :param node: The node to draw for :type node: :class:`aiida.orm.nodes.data.Data` :param node_label: The label to use for the nodes :type node_label: str :param show_pk: Show the PK of nodes alongside the label :type show_pk: bool :param dist: The number of steps away from this node to branch out :type dist: int :param follow_links_of_type: Follow links of this type when making steps, if None then it will follow CREATE and INPUT links :type follow_links_of_type: str """ import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('function is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member return get_ascii_tree(node, node_label, show_pk, dist, follow_links_of_type, True) def get_ascii_tree(node, node_label=None, show_pk=True, max_depth=1, follow_links_of_type=None, descend=True): """ Get a string representing an ASCII tree for the given node. .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. :param node: The node to get the tree for :type node: :class:`aiida.orm.nodes.node.Node` :param node_label: What to label the nodes with (can be an attribute name) :type node_label: str :param show_pk: If True, show the pk with the node label :type show_pk: bool :param max_depth: The maximum depth to follow starting from the node :type max_depth: int :param follow_links_of_type: Follow links of a given type, can be None :type follow_links_of_type: One of the members from :class:`aiida.common.links.LinkType` :param descend: if True will follow outputs, if False inputs :type descend: bool :return: The string giving an ASCII representation of the tree from the node :rtype: str """ import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('function is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member from ete3 import Tree tree_string = build_tree(node, node_label, show_pk, max_depth, follow_links_of_type, descend) tree = Tree(f'({tree_string});', format=1) return tree.get_ascii(show_internal=True) def build_tree(node, node_label=None, show_pk=True, max_depth=1, follow_links_of_type=None, descend=True, depth=0): """ Recursively build an ASCII string representation of the node tree .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. :param node: The node to get the tree for :type node: :class:`aiida.orm.nodes.node.Node` :param node_label: What to label the nodes with (can be an attribute name) :type node_label: str :param show_pk: If True, show the pk with the node label :type show_pk: bool :param max_depth: The maximum depth to follow starting from the node :type max_depth: int :param follow_links_of_type: Follow links of a given type, can be None :type follow_links_of_type: One of the members from :class:`aiida.common.links.LinkType` :param descend: if True will follow outputs, if False inputs :type descend: bool :param depth: the current depth :type depth: int :return: The string giving an ASCII representation of the tree from the node :rtype: str """ # pylint: disable=too-many-arguments import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('function is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member out_values = [] if depth < max_depth: relatives = [] if descend: outputs = node.get_outgoing(link_type=follow_links_of_type).all_nodes() else: # ascend if follow_links_of_type is None: follow_links_of_type = (LinkType.CREATE, LinkType.INPUT_CALC, LinkType.INPUT_WORK) outputs = node.get_incoming(link_type=follow_links_of_type).all_nodes() for child in sorted(outputs, key=lambda node: node.ctime): relatives.append( build_tree(child, node_label, show_pk, max_depth, follow_links_of_type, descend, depth + 1) ) if relatives: out_values.append(f"({', '.join(relatives)})") out_values.append(_generate_node_label(node, node_label, show_pk)) return ''.join(out_values) def _generate_node_label(node, node_attr, show_pk): """ Generate a label for the node. .. deprecated:: 1.1.0 Will be removed in `v2.0.0`. :param node: The node to generate the label for :type node: :class:`aiida.orm.nodes.node.Node` :param node_attr: The attribute to use as the label, can be None :type node_attr: str :param show_pk: if True, show the PK alongside the label :type show_pk: bool :return: The generated label :rtype: str """ import warnings from aiida.common.warnings import AiidaDeprecationWarning warnings.warn('function is deprecated and will be removed in `aiida-core==2.0.0`.', AiidaDeprecationWarning) # pylint: disable=no-member label = None if node_attr is None: try: label = node.process_label except AttributeError: label = None else: try: label = str(getattr(node, node_attr)) except AttributeError: try: label = node.get_attribute(node_attr) except AttributeError: pass # Couldn't find one, so just use the class name if label is None: label = node.__class__.__name__ if show_pk: label += f' [{node.pk}]' return label def calc_info(node): """Return a string with the summary of the state of a CalculationNode.""" from aiida.orm import ProcessNode, WorkChainNode if not isinstance(node, ProcessNode): raise TypeError(f'Unknown type: {type(node)}') process_label = node.process_label process_state = node.process_state.value.capitalize() exit_status = node.exit_status if exit_status is not None: string = f'{process_label}<{node.pk}> {process_state} [{exit_status}]' else: string = f'{process_label}<{node.pk}> {process_state}' if isinstance(node, WorkChainNode) and node.stepper_state_info: string += f' [{node.stepper_state_info}]' return string [docs]def format_call_graph(calc_node, info_fn=calc_info): """ Print a tree like the POSIX tree command for the calculation call graph :param calc_node: The calculation node :param info_fn: An optional function that takes the node and returns a string of information to be displayed for each node. """ call_tree = build_call_graph(calc_node, info_fn=info_fn) return format_tree_descending(call_tree) def build_call_graph(calc_node, info_fn=calc_info): """Build the call graph of a given node.""" info_string = info_fn(calc_node) called = calc_node.called called.sort(key=lambda x: x.ctime) if called: return info_string, [build_call_graph(child, info_fn) for child in called] return info_string def format_tree_descending(tree, prefix='', pos=-1): """Format a descending tree.""" # pylint: disable=too-many-branches text = [] if isinstance(tree, tuple): info = tree[0] else: info = tree if pos == -1: pre = '' elif pos == 0: pre = f'{prefix}{TREE_FIRST_ENTRY}' elif pos == 1: pre = f'{prefix}{TREE_MIDDLE_ENTRY}' else: pre = f'{prefix}{TREE_LAST_ENTRY}' text.append(f'{pre}{info}') if isinstance(tree, tuple): _, value = tree num_entries = len(value) if pos in [-1, 2]: new_prefix = f'{prefix} ' else: new_prefix = f'{prefix}│ ' for i, entry in enumerate(value): if i == num_entries - 1: pos = 2 elif i == 0: pos = 0 else: pos = 1 text.append(format_tree_descending(entry, new_prefix, pos)) return '\n'.join(text)