Tidy + type annotations

Author: MetRonnie

cylc/uiserver/data_store_mgr.py

@@ -37,6 +37,7 @@
 from functools import partial
 from pathlib import Path
 import time
+from typing import Optional
 
 from cylc.flow.id import Tokens
 from cylc.flow.network.server import PB_METHOD_MAP
@@ -361,12 +362,11 @@ def reconcile_update(self, topic, delta, w_id):
             except Exception as exc:
                 self.log.exception(exc)
 
-    async def entire_workflow_update(self, ids=None):
+    async def entire_workflow_update(self, ids: Optional[list] = None) -> None:
         """Update entire local data-store of workflow(s).
 
         Args:
-            ids (list): List of workflow external IDs.
-
+            ids: List of workflow external IDs.
 
         """
         if ids is None:

cylc/uiserver/handlers.py

@@ -18,7 +18,7 @@
 import json
 import getpass
 import socket
-from typing import Callable, Union
+from typing import TYPE_CHECKING, Callable, Union
 
 from graphene_tornado.tornado_graphql_handler import TornadoGraphQLHandler
 from graphql import get_default_backend
@@ -35,6 +35,9 @@
 from cylc.uiserver.authorise import Authorization, AuthorizationMiddleware
 from cylc.uiserver.websockets import authenticated as websockets_authenticated
 
+if TYPE_CHECKING:
+    from graphql.execution import ExecutionResult
+
 
 ME = getpass.getuser()
 
@@ -337,8 +340,8 @@ def context(self):
     def prepare(self):
         super().prepare()
 
-    @web.authenticated
-    async def execute(self, *args, **kwargs):
+    @web.authenticated  # type: ignore[arg-type]
+    async def execute(self, *args, **kwargs) -> 'ExecutionResult':
         # Use own backend, and TornadoGraphQLHandler already does validation.
         return await self.schema.execute(
             *args,

cylc/uiserver/resolvers.py

@@ -19,11 +19,22 @@
 import os
 from copy import deepcopy
 from subprocess import Popen, PIPE, DEVNULL
+from typing import (
+    TYPE_CHECKING, Any, Dict, Iterable, List, Union
+)
 
 from graphql.language.base import print_ast
 
-from cylc.flow.network.resolvers import BaseResolvers
 from cylc.flow.data_store_mgr import WORKFLOW
+from cylc.flow.network.resolvers import BaseResolvers
+
+
+if TYPE_CHECKING:
+    from logging import Logger
+    from graphql import ResolveInfo
+    from cylc.flow.data_store_mgr import DataStoreMgr
+    from cylc.flow.id import Tokens
+    from cylc.uiserver.workflows_mgr import WorkflowsManager
 
 
 # show traceback from cylc commands
@@ -180,33 +191,47 @@ async def play(cls, workflows, args, workflows_mgr, log):
 class Resolvers(BaseResolvers):
     """UI Server context GraphQL query and mutation resolvers."""
 
-    workflows_mgr = None
-
-    def __init__(self, data, log, **kwargs):
+    def __init__(
+        self,
+        data: 'DataStoreMgr',
+        log: 'Logger',
+        workflows_mgr: 'WorkflowsManager',
+        **kwargs
+    ):
         super().__init__(data)
         self.log = log
+        self.workflows_mgr = workflows_mgr
 
         # Set extra attributes
         for key, value in kwargs.items():
             if hasattr(self, key):
                 setattr(self, key, value)
 
     # Mutations
-    async def mutator(self, info, *m_args):
+    async def mutator(
+        self,
+        info: 'ResolveInfo',
+        command: str,
+        w_args: Dict[str, Any],
+        _kwargs: Dict[str, Any],
+        _meta: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
         """Mutate workflow."""
-        req_meta = {}
-        _, w_args, _, _ = m_args
-        req_meta['auth_user'] = info.context.get(
-            'current_user', 'unknown user')
+        req_meta = {
+            'auth_user': info.context.get(  # type: ignore[union-attr]
+                'current_user', 'unknown user'
+            )
+        }
         w_ids = [
             flow[WORKFLOW].id
             for flow in await self.get_workflows_data(w_args)]
         if not w_ids:
             return [{
                 'response': (False, 'No matching workflows')}]
         # Pass the request to the workflow GraphQL endpoints
-        _, variables, _, _ = info.context.get('graphql_params')
-
+        _, variables, _, _ = info.context.get(  # type: ignore[union-attr]
+            'graphql_params'
+        )
         # Create a modified request string,
         # containing only the current mutation/field.
         operation_ast = deepcopy(info.operation)
@@ -216,14 +241,20 @@ async def mutator(self, info, *m_args):
             'request_string': print_ast(operation_ast),
             'variables': variables,
         }
-        return await self.workflows_mgr.multi_request(
+        return await self.workflows_mgr.multi_request(  # type: ignore # TODO
             'graphql', w_ids, graphql_args, req_meta=req_meta
         )
 
-    async def service(self, info, *m_args):
+    async def service(
+        self,
+        info: 'ResolveInfo',
+        command: str,
+        workflows: Iterable['Tokens'],
+        kwargs: Dict[str, Any]
+    ) -> List[Union[bool, str]]:
         return await Services.play(
-            m_args[1]['workflows'],
-            m_args[2],
+            workflows,
+            kwargs,
             self.workflows_mgr,
             log=self.log
         )

cylc/uiserver/schema.py

@@ -20,6 +20,7 @@
 """
 
 from functools import partial
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import graphene
 from graphene.types.generic import GenericScalar
@@ -36,25 +37,32 @@
     sstrip,
 )
 
+if TYPE_CHECKING:
+    from graphql import ResolveInfo
+    from cylc.uiserver.resolvers import Resolvers
 
-async def mutator(root, info, command=None, workflows=None,
-                  exworkflows=None, **args):
+
+async def mutator(
+    root: Optional[Any],
+    info: 'ResolveInfo',
+    *,
+    command: str,
+    workflows: Optional[List[str]] = None,
+    **kwargs: Any
+):
     """Call the resolver method that act on the workflow service
     via the internal command queue."""
     if workflows is None:
         workflows = []
-    if exworkflows is None:
-        exworkflows = []
-    w_args = {
-        'workflows': [Tokens(w_id) for w_id in workflows],
-        'exworkflows': [Tokens(w_id) for w_id in exworkflows],
-    }
-    if args.get('args', False):
-        args.update(args.get('args', {}))
-        args.pop('args')
-
-    resolvers = info.context.get('resolvers')
-    res = await resolvers.service(info, command, w_args, args)
+    parsed_workflows = [Tokens(w_id) for w_id in workflows]
+    if kwargs.get('args', False):
+        kwargs.update(kwargs.get('args', {}))
+        kwargs.pop('args')
+
+    resolvers: 'Resolvers' = (
+        info.context.get('resolvers')  # type: ignore[union-attr]
+    )
+    res = await resolvers.service(info, command, parsed_workflows, kwargs)
     return GenericResponse(result=res)
 
 

cylc/uiserver/tests/conftest.py

@@ -34,14 +34,14 @@
     PbWorkflow,
     PbFamilyProxy,
 )
-from cylc.flow.network import ZMQSocketBase
+from cylc.flow.network.client import WorkflowRuntimeClient
 from cylc.flow.workflow_files import ContactFileFields as CFF
 
 from cylc.uiserver.data_store_mgr import DataStoreMgr
 from cylc.uiserver.workflows_mgr import WorkflowsManager
 
 
-class AsyncClientFixture(ZMQSocketBase):
+class AsyncClientFixture(WorkflowRuntimeClient):
     pattern = zmq.REQ
     host = ''
     port = 0

cylc/uiserver/tests/test_workflows_mgr.py

@@ -48,7 +48,7 @@ async def test_workflow_request_client_timeout(
     async_client.will_return(ClientTimeout)
     ctx, msg = await workflow_request(client=async_client, command='')
     assert not ctx
-    assert 'timeout' in msg.lower()
+    assert 'timeout' in msg.lower()  # type: ignore[attr-defined]
 
 
 @pytest.mark.asyncio

cylc/uiserver/workflows_mgr.py

@@ -26,14 +26,16 @@
 from contextlib import suppress
 from getpass import getuser
 import sys
+from typing import (
+    TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple
+)
 
 import zmq.asyncio
 
 from cylc.flow.id import Tokens
 from cylc.flow.exceptions import ClientError, ClientTimeout
-from cylc.flow.network import API
+from cylc.flow.network import API, MSG_TIMEOUT
 from cylc.flow.network.client import WorkflowRuntimeClient
-from cylc.flow.network import MSG_TIMEOUT
 from cylc.flow.network.scan import (
     api_version,
     contact_info,
@@ -42,32 +44,35 @@
 )
 from cylc.flow.workflow_files import ContactFileFields as CFF
 
+if TYPE_CHECKING:
+    from logging import Logger
+
+
 CLIENT_TIMEOUT = 2.0
 
 
 async def workflow_request(
-    client,
-    command,
-    args=None,
-    timeout=None,
-    req_context=None,
+    client: WorkflowRuntimeClient,
+    command: str,
+    args: Optional[Dict[str, Any]] = None,
+    timeout: Optional[float] = None,
+    req_context: Optional[str] = None,
     *,
-    log=None,
-    req_meta=None
-):
+    log: Optional['Logger'] = None,
+    req_meta: Optional[Dict[str, Any]] = None
+) -> Tuple[str, object]:
     """Workflow request command.
 
     Args:
-        client (WorkflowRuntimeClient): Instantiated workflow client.
-        command (str): Command/Endpoint name.
-        args (dict): Endpoint arguments.
-        timeout (float): Client request timeout (secs).
-        req_context (str): A string to identifier.
-        req_meta (dict): Meta data related to request, e.g. auth_user
+        client: Instantiated workflow client.
+        command: Command/Endpoint name.
+        args: Endpoint arguments.
+        timeout: Client request timeout (secs).
+        req_context: A string to identifier.
+        req_meta: Meta data related to request, e.g. auth_user
 
     Returns:
         tuple: (req_context, result)
-
     """
     if req_context is None:
         req_context = command
@@ -277,13 +282,13 @@ async def update(self):
 
     async def multi_request(
         self,
-        command,
-        workflows,
-        args=None,
-        multi_args=None,
+        command: str,
+        workflows: Iterable[str],
+        args: Optional[Dict[str, Any]] = None,
+        multi_args: Optional[Dict[str, Any]] = None,
         timeout=None,
-        req_meta=None
-    ):
+        req_meta: Optional[Dict[str, Any]] = None
+    ) -> List[object]:
         """Send requests to multiple workflows."""
         if args is None:
             args = {}