Add --group option to verdi process list (#2254)

This option allows a user to narrow down the query set to only those nodes
that are within the specified group.

aiida/backends/tests/cmdline/commands/test_process.py

@@ -113,21 +113,14 @@ class TestVerdiProcess(AiidaTestCase):
 
     TEST_TIMEOUT = 5.
 
-    def setUp(self):
-        super(TestVerdiProcess, self).setUp()
-        self.cli_runner = CliRunner()
-
-    def test_list(self):
-        """Test the list command."""
+    @classmethod
+    def setUpClass(cls, *args, **kwargs):
         # pylint: disable=protected-access
+        super(TestVerdiProcess, cls).setUpClass(*args, **kwargs)
+        from aiida.orm.groups import Group
         from aiida.work.processes import ProcessState
 
-        # Number of output lines in -r/--raw format should be zero when there are no calculations yet
-        result = self.cli_runner.invoke(cmd_process.process_list, ['-r'])
-        self.assertIsNone(result.exception, result.output)
-        self.assertEqual(len(get_result_lines(result)), 0)
-
-        calcs = []
+        cls.calcs = []
 
         # Create 6 WorkFunctionNodes and WorkChainNodes (one for each ProcessState)
         for state in ProcessState:
@@ -140,7 +133,7 @@ def test_list(self):
                 calc._set_exit_status(0)
 
             calc.store()
-            calcs.append(calc)
+            cls.calcs.append(calc)
 
             calc = WorkChainNode()
             calc._set_process_state(state)
@@ -150,7 +143,17 @@ def test_list(self):
                 calc._set_exit_status(1)
 
             calc.store()
-            calcs.append(calc)
+            cls.calcs.append(calc)
+
+        cls.group = Group('some_group').store()
+        cls.group.add_nodes(cls.calcs[0])
+
+    def setUp(self):
+        super(TestVerdiProcess, self).setUp()
+        self.cli_runner = CliRunner()
+
+    def test_list(self):
+        """Test the list command."""
 
         # Default behavior should yield all active states (CREATED, RUNNING and WAITING) so six in total
         result = self.cli_runner.invoke(cmd_process.process_list, ['-r'])
@@ -196,7 +199,14 @@ def test_list(self):
             self.assertEqual(len(get_result_lines(result)), 6)
 
             for line in get_result_lines(result):
-                self.assertIn(line.strip(), [str(calc.pk) for calc in calcs])
+                self.assertIn(line.strip(), [str(calc.pk) for calc in self.calcs])
+
+        # The group option should limit the query set to nodes in the group
+        for flag in ['-G', '--group']:
+            result = self.cli_runner.invoke(cmd_process.process_list, ['-r', '-P', 'pk', flag, str(self.group.pk)])
+            self.assertClickResultNoException(result)
+            self.assertEqual(len(get_result_lines(result)), 1)
+            self.assertEqual(get_result_lines(result)[0], str(self.calcs[0].pk))
 
     def test_process_show(self):
         """Test verdi process show"""

aiida/cmdline/commands/cmd_process.py

@@ -30,6 +30,7 @@ def verdi_process():
 @verdi_process.command('list')
 @options.PROJECT(
     type=click.Choice(CalculationQueryBuilder.valid_projections), default=CalculationQueryBuilder.default_projections)
+@options.GROUP(help='Only include entries that are a member of this group.')
 @options.ALL(help='Show all entries, regardless of their process state.')
 @options.PROCESS_STATE()
 @options.EXIT_STATUS()
@@ -38,14 +39,19 @@ def verdi_process():
 @options.LIMIT()
 @options.RAW()
 @decorators.with_dbenv()
-def process_list(all_entries, process_state, exit_status, failed, past_days, limit, project, raw):
+def process_list(all_entries, group, process_state, exit_status, failed, past_days, limit, project, raw):
     """Show a list of processes that are still running."""
     from tabulate import tabulate
     from aiida.cmdline.utils.common import print_last_process_state_change
 
+    relationships = {}
+
+    if group:
+        relationships['with_node'] = group
+
     builder = CalculationQueryBuilder()
     filters = builder.get_filters(all_entries, process_state, exit_status, failed)
-    query_set = builder.get_query_set(filters=filters, past_days=past_days, limit=limit)
+    query_set = builder.get_query_set(relationships=relationships, filters=filters, past_days=past_days, limit=limit)
     projected = builder.get_projected(query_set, projections=project)
 
     headers = projected.pop(0)

aiida/cmdline/utils/query/calculation.py

@@ -77,10 +77,13 @@ def get_filters(self, all_entries=False, process_state=None, exit_status=None, f
 
         return filters
 
-    def get_query_set(self, filters=None, order_by=None, past_days=None, limit=None):
+    def get_query_set(self, relationships=None, filters=None, order_by=None, past_days=None, limit=None):
         """
         Return the query set of calculations for the given filters and query parameters
 
+        :param relationships: a mapping of relationships to join on, e.g. {'with_node': Group} to join on a Group. The
+            keys in this dictionary should be the keyword used in the `append` method of the `QueryBuilder` to join the
+            entity on that is defined as the value.
         :param filters: rules to filter query results with
         :param order_by: order the query set by this criterion
         :param past_days: only include entries from the last past days
@@ -104,6 +107,10 @@ def get_query_set(self, filters=None, order_by=None, past_days=None, limit=None)
         builder = QueryBuilder()
         builder.append(cls=ProcessNode, filters=filters, project=projected_attributes, tag='process')
 
+        if relationships is not None:
+            for tag, entity in relationships.items():
+                builder.append(cls=type(entity), filters={'id': entity.id}, **{tag: 'process'})
+
         if order_by is not None:
             builder.order_by({'process': order_by})
         else: