Source code for aiida.cmdline.utils.ascii_vis

# -*- 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."""
__all__ = ('format_call_graph',)

TREE_LAST_ENTRY = '\u2514\u2500\u2500 '
TREE_MIDDLE_ENTRY = '\u251C\u2500\u2500 '
TREE_FIRST_ENTRY = TREE_MIDDLE_ENTRY


[docs]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)
[docs]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
[docs]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)