Start pattern to disallow untyped functions, MyPY updates (#91)

* Start on enforcing some type hints

Add required type hints for control module

* Make control module mock dispatcher conform to the main protocol

Author: AlanCoding

.github/workflows/ci.yml

@@ -62,7 +62,7 @@ jobs:
       - run: pip install mypy
       - run: pip install -e .[pg_notify]
       - run: python3 -m pip install types-PyYAML
-      - run: mypy --ignore-missing-imports dispatcher
+      - run: mypy dispatcher
 
   flake8:
     name: Run flake8

Makefile

@@ -19,4 +19,4 @@ linters:
 	black dispatcher/
 	isort dispatcher/
 	flake8 dispatcher/
-	mypy --ignore-missing-imports dispatcher
+	mypy dispatcher

dispatcher/control.py

@@ -3,10 +3,11 @@
 import logging
 import time
 import uuid
-from typing import Optional
+from typing import Optional, Union
 
 from .factories import get_broker
 from .producers import BrokeredProducer
+from .protocols import Producer
 from .service.asyncio_tasks import ensure_fatal
 
 logger = logging.getLogger('awx.main.dispatch.control')
@@ -23,49 +24,57 @@ class ControlCallbacks:
     it exists to interact with producers, using variables relevant to the particular
     control message being sent"""
 
-    def __init__(self, queuename, send_data, expected_replies) -> None:
+    def __init__(self, queuename: Optional[str], send_data: dict, expected_replies: int) -> None:
         self.queuename = queuename
         self.send_data = send_data
         self.expected_replies = expected_replies
 
         # received_replies only tracks the reply message, not the channel name
         # because they come via a temporary reply_to channel and that is not user-facing
-        self.received_replies: list[str] = []
+        self.received_replies: list[Union[dict, str]] = []
         self.events = ControlEvents()
         self.shutting_down = False
 
-    async def process_message(self, payload, producer=None, channel=None) -> tuple[Optional[str], Optional[str]]:
+    async def process_message(
+        self, payload: Union[dict, str], producer: Optional[Producer] = None, channel: Optional[str] = None
+    ) -> tuple[Optional[str], Optional[str]]:
         self.received_replies.append(payload)
         if self.expected_replies and (len(self.received_replies) >= self.expected_replies):
             self.events.exit_event.set()
         return (None, None)
 
-    async def connected_callback(self, producer) -> None:
+    async def connected_callback(self, producer: Producer) -> None:
         payload = json.dumps(self.send_data)
-        await producer.notify(channel=self.queuename, message=payload)
+        # Ignore the type hint here because we know it is a brokered producer
+        await producer.notify(channel=self.queuename, message=payload)  # type: ignore[attr-defined]
         logger.info('Sent control message, expecting replies soon')
 
+    async def main(self) -> None:
+        "Unused"
+        pass
+
 
 class Control(object):
     def __init__(self, broker_name: str, broker_config: dict, queue: Optional[str] = None) -> None:
         self.queuename = queue
         self.broker_name = broker_name
         self.broker_config = broker_config
 
-    def running(self, *args, **kwargs):
-        return self.control_with_reply('running', *args, **kwargs)
-
-    def cancel(self, task_ids, with_reply=True):
-        if with_reply:
-            return self.control_with_reply('cancel', extra_data={'task_ids': task_ids})
-        else:
-            self.control({'control': 'cancel', 'task_ids': task_ids, 'reply_to': None}, extra_data={'task_ids': task_ids})
-
     @classmethod
-    def generate_reply_queue_name(cls):
+    def generate_reply_queue_name(cls) -> str:
         return f"reply_to_{str(uuid.uuid4()).replace('-', '_')}"
 
-    async def acontrol_with_reply_internal(self, producer, send_data, expected_replies, timeout):
+    @staticmethod
+    def parse_replies(received_replies: list[Union[str, dict]]) -> list[dict]:
+        ret = []
+        for payload in received_replies:
+            if isinstance(payload, dict):
+                ret.append(payload)
+            else:
+                ret.append(json.loads(payload))
+        return ret
+
+    async def acontrol_with_reply_internal(self, producer: Producer, send_data: dict, expected_replies: int, timeout: float) -> list[dict]:
         control_callbacks = ControlCallbacks(self.queuename, send_data, expected_replies)
 
         await producer.start_producing(control_callbacks)
@@ -83,22 +92,22 @@ async def acontrol_with_reply_internal(self, producer, send_data, expected_repli
         control_callbacks.shutting_down = True
         await producer.shutdown()
 
-        return [json.loads(payload) for payload in control_callbacks.received_replies]
+        return self.parse_replies(control_callbacks.received_replies)
 
-    def make_producer(self, reply_queue):
+    def make_producer(self, reply_queue: str) -> Producer:
         broker = get_broker(self.broker_name, self.broker_config, channels=[reply_queue])
         return BrokeredProducer(broker, close_on_exit=True)
 
-    async def acontrol_with_reply(self, command, expected_replies=1, timeout=1, data=None):
+    async def acontrol_with_reply(self, command: str, expected_replies: int = 1, timeout: int = 1, data: Optional[dict] = None) -> list[dict]:
         reply_queue = Control.generate_reply_queue_name()
-        send_data = {'control': command, 'reply_to': reply_queue}
+        send_data: dict[str, Union[dict, str]] = {'control': command, 'reply_to': reply_queue}
         if data:
             send_data['control_data'] = data
 
         return await self.acontrol_with_reply_internal(self.make_producer(reply_queue), send_data, expected_replies, timeout)
 
-    async def acontrol(self, command, data=None):
-        send_data = {'control': command}
+    async def acontrol(self, command: str, data: Optional[dict] = None) -> None:
+        send_data: dict[str, Union[dict, str]] = {'control': command}
         if data:
             send_data['control_data'] = data
 
@@ -110,13 +119,13 @@ def control_with_reply(self, command: str, expected_replies: int = 1, timeout: f
         logger.info('control-and-reply {} to {}'.format(command, self.queuename))
         start = time.time()
         reply_queue = Control.generate_reply_queue_name()
-        send_data = {'control': command, 'reply_to': reply_queue}
+        send_data: dict[str, Union[dict, str]] = {'control': command, 'reply_to': reply_queue}
         if data:
             send_data['control_data'] = data
 
         broker = get_broker(self.broker_name, self.broker_config, channels=[reply_queue])
 
-        def connected_callback():
+        def connected_callback() -> None:
             payload = json.dumps(send_data)
             if self.queuename:
                 broker.publish_message(channel=self.queuename, message=payload)
@@ -131,9 +140,9 @@ def connected_callback():
         logger.info(f'control-and-reply message returned in {time.time() - start} seconds')
         return replies
 
-    def control(self, command, data=None):
+    def control(self, command: str, data: Optional[dict] = None) -> None:
         "Send message in fire-and-forget mode, as synchronous code. Only for no-reply control."
-        send_data = {'control': command}
+        send_data: dict[str, Union[dict, str]] = {'control': command}
         if data:
             send_data['control_data'] = data
 

dispatcher/protocols.py

@@ -65,7 +65,6 @@ async def dispatch_task(self, message: dict) -> None:
 
 
 class DispatcherMain(Protocol):
-    fd_lock: asyncio.Lock
 
     async def main(self) -> None:
         """This is the method that runs the service, bring your own event loop"""

mypy.ini

@@ -0,0 +1,4 @@
+[mypy]
+
+[mypy-dispatcher.control]
+disallow_untyped_defs = True