From 183c34e7ee6ad567e7a8637ea18b6265d84e8494 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Thu, 4 Jul 2024 15:25:13 +0200 Subject: [PATCH] CLI: Accept mulitple node identifiers in `verdi node graph generate` (#6443) The `--identifier` option allows the user to specify which identifier to label nodes in the graph with: pk, uuid or label. Here, the interface is updated to allow specifying multiple identifiers, e.g.: verdi node graph generate --identifier pk uuid -- If more than one identifier type is specified, the resulting identifiers for each node are joined using a `|` character. --- src/aiida/cmdline/commands/cmd_node.py | 9 +++-- src/aiida/tools/visualization/graph.py | 30 ++++++++------ tests/cmdline/commands/test_node.py | 2 +- tests/tools/visualization/test_graph.py | 39 +++++++++++++++++-- .../test_graph_node_identifiers_label_.txt | 9 +++++ ..._graph_node_identifiers_node_id_type3_.txt | 9 +++++ ..._graph_node_identifiers_node_id_type4_.txt | 9 +++++ .../test_graph_node_identifiers_pk_.txt | 9 +++++ .../test_graph_node_identifiers_uuid_.txt | 9 +++++ 9 files changed, 106 insertions(+), 19 deletions(-) create mode 100644 tests/tools/visualization/test_graph/test_graph_node_identifiers_label_.txt create mode 100644 tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type3_.txt create mode 100644 tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type4_.txt create mode 100644 tests/tools/visualization/test_graph/test_graph_node_identifiers_pk_.txt create mode 100644 tests/tools/visualization/test_graph/test_graph_node_identifiers_uuid_.txt diff --git a/src/aiida/cmdline/commands/cmd_node.py b/src/aiida/cmdline/commands/cmd_node.py index 8e6ae8fba0..79efcebcef 100644 --- a/src/aiida/cmdline/commands/cmd_node.py +++ b/src/aiida/cmdline/commands/cmd_node.py @@ -15,6 +15,7 @@ from aiida.cmdline.commands.cmd_verdi import verdi from aiida.cmdline.params import arguments, options +from aiida.cmdline.params.options.multivalue import MultipleValueOption from aiida.cmdline.params.types.plugin import PluginParamType from aiida.cmdline.utils import decorators, echo, echo_tabulate, multi_line_input from aiida.cmdline.utils.decorators import with_dbenv @@ -439,8 +440,10 @@ def verdi_graph(): ) @click.option( '--identifier', + 'identifiers', help='the type of identifier to use within the node text', - default='uuid', + default=('uuid',), + cls=MultipleValueOption, type=click.Choice(['pk', 'uuid', 'label']), ) @click.option( @@ -483,7 +486,7 @@ def verdi_graph(): def graph_generate( root_nodes, link_types, - identifier, + identifiers, ancestor_depth, descendant_depth, process_out, @@ -506,7 +509,7 @@ def graph_generate( output_file = pathlib.Path(f'{pks}.{engine}.{output_format}') echo.echo_info(f'Initiating graphviz engine: {engine}') - graph = Graph(engine=engine, node_id_type=identifier) + graph = Graph(engine=engine, node_id_type=identifiers) link_types = {'all': (), 'logic': ('input_work', 'return'), 'data': ('input_calc', 'create')}[link_types] for root_node in root_nodes: diff --git a/src/aiida/tools/visualization/graph.py b/src/aiida/tools/visualization/graph.py index 2fe7f4250e..91411796ea 100644 --- a/src/aiida/tools/visualization/graph.py +++ b/src/aiida/tools/visualization/graph.py @@ -29,6 +29,7 @@ __all__ = ('Graph', 'default_link_styles', 'default_node_styles', 'pstate_node_styles', 'default_node_sublabels') LinkAnnotateType = Literal[None, 'label', 'type', 'both'] +IdentifierType = Literal['pk', 'uuid', 'label'] class LinkStyleFunc(Protocol): @@ -254,18 +255,25 @@ def default_node_sublabels(node: orm.Node) -> str: return sublabel -def get_node_id_label(node: orm.Node, id_type: Literal['pk', 'uuid', 'label']) -> str: +NODE_IDENTIFIER_TO_LABEL = { + 'pk': lambda node: str(node.pk), + 'uuid': lambda node: node.uuid.split('-')[0], + 'label': lambda node: node.label, +} + + +def get_node_id_label(node: orm.Node, id_type: IdentifierType | list[IdentifierType]) -> str: """Return an identifier str for the node""" - if id_type == 'pk': - return str(node.pk) - if id_type == 'uuid': - return node.uuid.split('-')[0] - if id_type == 'label': - return node.label - raise ValueError(f'node_id_type not recognised: {id_type}') + + id_types = id_type if isinstance(id_type, (list, tuple)) else [id_type] + + try: + return '|'.join(NODE_IDENTIFIER_TO_LABEL[key](node) for key in id_types) + except KeyError as exception: + raise ValueError(f'`{id_type}` is not a valid `node_id_type`, choose from: pk, uuid, label') from exception -def _get_node_label(node: orm.Node, id_type: Literal['pk', 'uuid', 'label'] = 'pk') -> str: +def _get_node_label(node: orm.Node, id_type: IdentifierType | list[IdentifierType] = 'pk') -> str: """Return a label text of node and the return format is ' ()'.""" if isinstance(node, orm.Data): label = f'{node.__class__.__name__} ({get_node_id_label(node, id_type)})' @@ -287,7 +295,7 @@ def _add_graphviz_node( node_sublabel_func, style_override: None | dict = None, include_sublabels: bool = True, - id_type: Literal['pk', 'uuid', 'label'] = 'pk', + id_type: IdentifierType | list[IdentifierType] = 'pk', ): """Create a node in the graph @@ -360,7 +368,7 @@ def __init__( link_style_fn: LinkStyleFunc | None = None, node_style_fn: Callable[[orm.Node], dict] | None = None, node_sublabel_fn: Callable[[orm.Node], str] | None = None, - node_id_type: Literal['pk', 'uuid', 'label'] = 'pk', + node_id_type: IdentifierType | list[IdentifierType] = 'pk', backend: StorageBackend | None = None, ): """A class to create graphviz graphs of the AiiDA node provenance diff --git a/tests/cmdline/commands/test_node.py b/tests/cmdline/commands/test_node.py index d9c7774556..66ca83b686 100644 --- a/tests/cmdline/commands/test_node.py +++ b/tests/cmdline/commands/test_node.py @@ -392,7 +392,7 @@ def test_node_id_label_format(self, run_cli_command): filename = f'{root_node}.dot.pdf' for id_label_type in ['uuid', 'pk', 'label']: - options = ['--identifier', id_label_type, root_node] + options = ['--identifier', id_label_type, '--', root_node] try: run_cli_command(cmd_node.graph_generate, options) assert os.path.isfile(filename) diff --git a/tests/tools/visualization/test_graph.py b/tests/tools/visualization/test_graph.py index 17bda9993f..ba47b335b2 100644 --- a/tests/tools/visualization/test_graph.py +++ b/tests/tools/visualization/test_graph.py @@ -8,6 +8,8 @@ ########################################################################### """Tests for creating graphs (using graphviz)""" +import re + import pytest from aiida import orm from aiida.common import AttributeDict @@ -290,10 +292,6 @@ def test_graph_graphviz_source_pstate(self): graph = graph_mod.Graph(node_style_fn=graph_mod.pstate_node_styles) graph.recurse_descendants(nodes.pd0) - # print() - # print(graph.graphviz.source) - # graph.graphviz.render("test_graphviz_pstate", cleanup=True) - expected = """\ digraph {{ N{pd0} [label="Dict ({pd0})" color=red pencolor=black penwidth=6 shape=rectangle] @@ -325,3 +323,36 @@ def test_graph_graphviz_source_pstate(self): assert sorted([line.strip() for line in graph.graphviz.source.splitlines()]) == sorted( [line.strip() for line in expected.splitlines()] ) + + @pytest.mark.parametrize( + 'node_id_type', + ( + 'pk', + 'uuid', + 'label', + ('pk', 'uuid'), + ('pk', 'label'), + ), + ) + def test_graph_node_identifiers(self, node_id_type, monkeypatch, file_regression): + """.""" + nodes = self.create_provenance() + + # Monkeypatch the mapping of lambdas that convert return a node's identifier in string form. This is because + # the pks and uuids of the test nodes will change between each test run and this would fail the file regression. + node_identifier_to_label = { + 'pk': lambda node: '10', + 'uuid': lambda node: '16739459', + 'label': lambda node: 'some-label', + } + monkeypatch.setattr(graph_mod, 'NODE_IDENTIFIER_TO_LABEL', node_identifier_to_label) + + graph = graph_mod.Graph(node_id_type=node_id_type) + graph.recurse_descendants(nodes.calcf1) + + # The order of certain output lines can be randomly ordered so we split the file in lines, sort, and then join + # them into a single string again. The node identifiers generated by the engine are of the form ``N{pk}`` and + # they also clearly vary, so they are replaced with the ``NODE`` placeholder. + string = '\n'.join(sorted(graph.graphviz.source.strip().split('\n'))) + string = re.sub(r'N\d+', 'NODE', string) + file_regression.check(string) diff --git a/tests/tools/visualization/test_graph/test_graph_node_identifiers_label_.txt b/tests/tools/visualization/test_graph/test_graph_node_identifiers_label_.txt new file mode 100644 index 0000000000..2c23527560 --- /dev/null +++ b/tests/tools/visualization/test_graph/test_graph_node_identifiers_label_.txt @@ -0,0 +1,9 @@ + NODE -> NODE [color="#000000" style=solid] + NODE -> NODE [color="#000000" style=solid] + NODE [label="CalcFunctionNode (some-label) + NODE [label="Dict (some-label)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] + NODE [label="FolderData (some-label)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] +Exit Code: 200" color=red fillcolor="#de707f77" penwidth=6 shape=rectangle style=filled] +State: finished +digraph { +} \ No newline at end of file diff --git a/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type3_.txt b/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type3_.txt new file mode 100644 index 0000000000..7644ab1562 --- /dev/null +++ b/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type3_.txt @@ -0,0 +1,9 @@ + NODE -> NODE [color="#000000" style=solid] + NODE -> NODE [color="#000000" style=solid] + NODE [label="CalcFunctionNode (10|16739459) + NODE [label="Dict (10|16739459)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] + NODE [label="FolderData (10|16739459)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] +Exit Code: 200" color=red fillcolor="#de707f77" penwidth=6 shape=rectangle style=filled] +State: finished +digraph { +} \ No newline at end of file diff --git a/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type4_.txt b/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type4_.txt new file mode 100644 index 0000000000..87b15dcf5c --- /dev/null +++ b/tests/tools/visualization/test_graph/test_graph_node_identifiers_node_id_type4_.txt @@ -0,0 +1,9 @@ + NODE -> NODE [color="#000000" style=solid] + NODE -> NODE [color="#000000" style=solid] + NODE [label="CalcFunctionNode (10|some-label) + NODE [label="Dict (10|some-label)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] + NODE [label="FolderData (10|some-label)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] +Exit Code: 200" color=red fillcolor="#de707f77" penwidth=6 shape=rectangle style=filled] +State: finished +digraph { +} \ No newline at end of file diff --git a/tests/tools/visualization/test_graph/test_graph_node_identifiers_pk_.txt b/tests/tools/visualization/test_graph/test_graph_node_identifiers_pk_.txt new file mode 100644 index 0000000000..ec6b88d6c2 --- /dev/null +++ b/tests/tools/visualization/test_graph/test_graph_node_identifiers_pk_.txt @@ -0,0 +1,9 @@ + NODE -> NODE [color="#000000" style=solid] + NODE -> NODE [color="#000000" style=solid] + NODE [label="CalcFunctionNode (10) + NODE [label="Dict (10)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] + NODE [label="FolderData (10)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] +Exit Code: 200" color=red fillcolor="#de707f77" penwidth=6 shape=rectangle style=filled] +State: finished +digraph { +} \ No newline at end of file diff --git a/tests/tools/visualization/test_graph/test_graph_node_identifiers_uuid_.txt b/tests/tools/visualization/test_graph/test_graph_node_identifiers_uuid_.txt new file mode 100644 index 0000000000..b4ca47ff60 --- /dev/null +++ b/tests/tools/visualization/test_graph/test_graph_node_identifiers_uuid_.txt @@ -0,0 +1,9 @@ + NODE -> NODE [color="#000000" style=solid] + NODE -> NODE [color="#000000" style=solid] + NODE [label="CalcFunctionNode (16739459) + NODE [label="Dict (16739459)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] + NODE [label="FolderData (16739459)" fillcolor="#8cd499ff" penwidth=0 shape=ellipse style=filled] +Exit Code: 200" color=red fillcolor="#de707f77" penwidth=6 shape=rectangle style=filled] +State: finished +digraph { +} \ No newline at end of file