completion_server: support "cylc set" arguments

cylc/flow/scripts/completion_server.py

@@ -40,6 +40,8 @@
 #   Which provide possible values to the completion functions.
 
 import asyncio
+from contextlib import suppress
+import inspect
 import os
 from pathlib import Path
 import select
@@ -50,6 +52,7 @@
 from packaging.specifiers import SpecifierSet
 
 from cylc.flow.cfgspec.glbl_cfg import glbl_cfg
+from cylc.flow.exceptions import CylcError
 from cylc.flow.id import tokenise, IDTokens, Tokens
 from cylc.flow.network.scan import scan
 from cylc.flow.option_parsers import CylcOptionParser as COP
@@ -193,7 +196,12 @@ async def complete_cylc(_root: str, *items: str) -> t.List[str]:
         if ret is not None:
             return ret
     if previous and previous.startswith('-'):
-        ret = await complete_option_value(command, previous, partial)
+        ret = await complete_option_value(
+            command,
+            previous,
+            partial,
+            items=items,
+        )
         if ret is not None:
             return ret
 
@@ -256,10 +264,11 @@ async def complete_option(
 async def complete_option_value(
     command: str,
     option: str,
-    partial: t.Optional[str] = None
+    partial: t.Optional[str] = None,
+    items: t.Optional[t.Iterable[str]] = None,
 ) -> t.Optional[t.List[str]]:
     """Complete values for --options."""
-    vals = await list_option_values(command, option, partial)
+    vals = await list_option_values(command, option, partial, items=items)
     if vals is not None:
         return complete(partial, vals)
     return None
@@ -331,17 +340,44 @@ async def list_option_values(
     command: str,
     option: str,
     partial: t.Optional[str] = '',
+    items: t.Optional[t.Iterable[str]] = None,
 ) -> t.Optional[t.List[str]]:
     """List values for an option in a Cylc command.
 
+    Args:
+        command:
+            The Cylc sub-command.
+        option:
+            The --option to list possible values for.
+        partial:
+            The part of the command the user is completing.
+        items:
+            The CLI context, i.e. everything that has been typed on the CLI
+            before the partial.
+
     E.G. --flow ['all', 'new', 'none']
     """
     if option in OPTION_MAP:
         list_option = OPTION_MAP[option]
         if not list_option:
             # do not perform completion for this option
             return []
-        return await list_option(None, partial)
+        kwargs = {}
+        if 'tokens_list' in inspect.getfullargspec(list_option).args:
+            # the function requires information about tokens already specified
+            # on the CLI
+            # (e.g. the workflow//cycle/task the command is operating on)
+            tokens_list = []
+            for item in items or []:
+                # pull out things from the command which look like IDs
+                if '//' in item:
+                    with suppress(ValueError):
+                        tokens_list.append(Tokens(item))
+                        continue
+                    with suppress(ValueError):
+                        tokens_list.append(Tokens(item, relative=True))
+            kwargs['tokens_list'] = tokens_list
+        return await list_option(partial, **kwargs)
     return None
 
 
@@ -413,7 +449,6 @@ async def list_resources(_partial: str) -> t.List[str]:
 
 
 async def list_dir(
-    _workflow: t.Optional[str],
     partial: t.Optional[str]
 ) -> t.List[str]:
     """List an arbitrary dir on the filesystem.
@@ -460,21 +495,103 @@ def list_rel_dir(path: Path, base: Path) -> t.List[str]:
 
 
 async def list_flows(
-    _workflow: t.Optional[str],
     _partial: t.Optional[str]
 ) -> t.List[str]:
     """List values for the --flow option."""
     return ['all', 'none', 'new']
 
 
 async def list_colours(
-    _workflow: t.Optional[str],
     _partial: t.Optional[str]
 ) -> t.List[str]:
     """List values for the --color option."""
     return ['never', 'auto', 'always']
 
 
+async def list_outputs(
+    _partial: t.Optional[str],
+    tokens_list: t.Optional[t.List[Tokens]],
+):
+    """List task outputs."""
+    return (await _list_prereqs_and_outputs(tokens_list))[1]
+
+
+async def list_prereqs(
+    _partial: t.Optional[str],
+    tokens_list: t.Optional[t.List[Tokens]],
+):
+    """List task prerequisites."""
+    return (await _list_prereqs_and_outputs(tokens_list))[0] + ['all']
+
+
+async def _list_prereqs_and_outputs(
+    tokens_list: t.Optional[t.List[Tokens]],
+) -> t.Tuple[t.List[str], t.List[str]]:
+    """List task prerequisites and outputs.
+
+    Returns:
+        tuple - (prereqs, outputs)
+
+    """
+    if not tokens_list:
+        # no context information available on the CLI
+        # we can't list prereqs/outputs
+        return ([], [])
+
+    # dynamic import for this relatively unlikely case to avoid slowing down
+    # server startup unnecessarily
+    from cylc.flow.network.client_factory import get_client
+    from cylc.flow.scripts.show import prereqs_and_outputs_query
+    from types import SimpleNamespace
+
+    workflows: t.Dict[str, t.List[Tokens]] = {}
+    current_workflow = None
+    for tokens in tokens_list:
+        workflow = tokens['workflow']
+        task = tokens['task']
+        if workflow:
+            workflows.setdefault(workflow, [])
+            current_workflow = workflow
+        if current_workflow and task:
+            workflows[current_workflow].append(tokens.task)
+
+    clients = {}
+    for workflow in workflows:
+        with suppress(CylcError):
+            clients[workflow] = get_client(workflow)
+
+    if not workflows:
+        return ([], [])
+
+    json: dict = {}
+    await asyncio.gather(*(
+        prereqs_and_outputs_query(
+            workflow,
+            workflows[workflow],
+            pclient,
+            SimpleNamespace(json=True),
+            json,
+        )
+        for workflow, pclient in clients.items()
+    ))
+
+    if not json:
+        return ([], [])
+    return (
+        [
+            f"{cond['taskId']}:{cond['reqState']}"
+            for value in json.values()
+            for prerequisite in value['prerequisites']
+            for cond in prerequisite['conditions']
+        ],
+        [
+            output['label']
+            for value in json.values()
+            for output in value['outputs']
+        ],
+    )
+
+
 # non-exhaustive list of Cylc commands which take non-workflow arguments
 COMMAND_MAP: t.Dict[str, t.Optional[t.Callable]] = {
     # register commands which have special positional arguments
@@ -513,6 +630,8 @@ async def list_colours(
     '--flow': list_flows,
     '--colour': list_colours,
     '--color': list_colours,
+    '--out': list_outputs,
+    '--pre': list_prereqs,
     # options for which we should not attempt to complete values for
     '--rm': None,
     '--run-name': None,
@@ -528,17 +647,21 @@ async def list_colours(
 }
 
 
-def cli_detokenise(tokens: Tokens) -> str:
+def cli_detokenise(tokens: Tokens, relative=False) -> str:
     """Format tokens for use on the command line.
 
     I.E. add the trailing slash[es] onto the end.
     """
     if tokens.is_null:
         # shouldn't happen but prevents possible error
         return ''
+    if relative:
+        id_ = tokens.relative_id
+    else:
+        id_ = tokens.id
     if tokens.lowest_token == IDTokens.Workflow.value:
-        return f'{tokens.id}//'
-    return f'{tokens.id}/'
+        return f'{id_}//'
+    return f'{id_}/'
 
 
 def next_token(tokens: Tokens) -> t.Optional[str]: