_frontend/cli.py: Add option to close multiple workspaces

At present, it is only possible to close workspaces for elements one at
a time. This can become slightly tedious process when you have multiple
workspaces open and you want to close all of them, maybe because you
just finished working on a set of related elements.

Instead of accepting a single element, accept a list of elements as
argument for `bst workspace close`. Additionally, add `-a`/`--all`
option to close all workspaces.

Fixes #337 - Add option to close all workspaces.

2 files changed, 56 insertions, 7 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index 19e78be4b..345e2eac7 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -579,22 +579,31 @@ def workspace_open(app, no_checkout, force, track_, element, directory):
 ##################################################################
 #                     Workspace Close Command                    #
 ##################################################################
-@workspace.command(name='close', short_help="Close a workspace")
+@workspace.command(name='close', short_help="Close workspaces")
 @click.option('--remove-dir', default=False, is_flag=True,
               help="Remove the path that contains the closed workspace")
-@click.argument('element',
+@click.option('--all', '-a', 'all_', default=False, is_flag=True,
+              help="Close all open workspaces")
+@click.argument('elements', nargs=-1,
                 type=click.Path(dir_okay=False, readable=True))
 @click.pass_obj
-def workspace_close(app, remove_dir, element):
+def workspace_close(app, remove_dir, all_, elements):
     """Close a workspace"""
 
+    if not (all_ or elements):
+        click.echo('ERROR: no elements specified', err=True)
+        sys.exit(-1)
+
     if app.interactive and remove_dir:
         if not click.confirm('This will remove all your changes, are you sure?'):
             click.echo('Aborting', err=True)
             sys.exit(-1)
 
     with app.partially_initialized():
-        app.close_workspace(element, remove_dir)
+        if all_:
+            elements = [element_name for element_name, _ in app.project.workspaces.list()]
+        for element in elements:
+            app.close_workspace(element, remove_dir)
 
 
 ##################################################################
diff --git a/tests/frontend/workspace.py b/tests/frontend/workspace.py
index 2dc8d2ab5..c1e11979a 100644
--- a/tests/frontend/workspace.py
+++ b/tests/frontend/workspace.py
@@ -17,12 +17,12 @@ DATA_DIR = os.path.join(
 )
 
 
-def open_workspace(cli, tmpdir, datafiles, kind, track):
+def open_workspace(cli, tmpdir, datafiles, kind, track, suffix=''):
     project = os.path.join(datafiles.dirname, datafiles.basename)
     bin_files_path = os.path.join(project, 'files', 'bin-files')
     element_path = os.path.join(project, 'elements')
-    element_name = 'workspace-test-{}.bst'.format(kind)
-    workspace = os.path.join(str(tmpdir), 'workspace')
+    element_name = 'workspace-test-{}{}.bst'.format(kind, suffix)
+    workspace = os.path.join(str(tmpdir), 'workspace{}'.format(suffix))
 
     # Create our repo object of the given source type with
     # the bin files, and then collect the initial ref.
@@ -154,6 +154,46 @@ def test_close_nonexistant_element(cli, tmpdir, datafiles):
 
 
 @pytest.mark.datafiles(DATA_DIR)
+def test_close_multiple(cli, tmpdir, datafiles):
+    tmpdir_alpha = os.path.join(str(tmpdir), 'alpha')
+    tmpdir_beta = os.path.join(str(tmpdir), 'beta')
+    alpha, project, workspace_alpha = open_workspace(
+        cli, tmpdir_alpha, datafiles, 'git', False, suffix='-alpha')
+    beta, project, workspace_beta = open_workspace(
+        cli, tmpdir_beta, datafiles, 'git', False, suffix='-beta')
+
+    # Close the workspaces
+    result = cli.run(project=project, args=[
+        'workspace', 'close', '--remove-dir', alpha, beta
+    ])
+    result.assert_success()
+
+    # Assert the workspace dirs have been deleted
+    assert not os.path.exists(workspace_alpha)
+    assert not os.path.exists(workspace_beta)
+
+
+@pytest.mark.datafiles(DATA_DIR)
+def test_close_all(cli, tmpdir, datafiles):
+    tmpdir_alpha = os.path.join(str(tmpdir), 'alpha')
+    tmpdir_beta = os.path.join(str(tmpdir), 'beta')
+    alpha, project, workspace_alpha = open_workspace(
+        cli, tmpdir_alpha, datafiles, 'git', False, suffix='-alpha')
+    beta, project, workspace_beta = open_workspace(
+        cli, tmpdir_beta, datafiles, 'git', False, suffix='-beta')
+
+    # Close the workspaces
+    result = cli.run(project=project, args=[
+        'workspace', 'close', '--remove-dir', '--all'
+    ])
+    result.assert_success()
+
+    # Assert the workspace dirs have been deleted
+    assert not os.path.exists(workspace_alpha)
+    assert not os.path.exists(workspace_beta)
+
+
+@pytest.mark.datafiles(DATA_DIR)
 def test_reset(cli, tmpdir, datafiles):
     # Open the workspace
     element_name, project, workspace = open_workspace(cli, tmpdir, datafiles, 'git', False)