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 -- <PK>

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

@@ -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:

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 '<NodeType> (<id>)'."""
     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

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)

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)

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 {
+}

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 {
+}

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 {
+}

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 {
+}

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 {
+}