Typing

worker_init_elements/worker.py

@@ -3,6 +3,7 @@ import sqlite3
 import sys
 import uuid
 from collections import OrderedDict
+from collections.abc import Iterator
 from logging import Logger, getLogger
 from time import sleep
 
@@ -25,7 +26,7 @@ PENDING_STATE = "pending"
 ERROR_STATE = "error"
 
 
-def split_chunks(items, n):
+def split_chunks(items: list, n: int) -> Iterator[list]:
     """
     Yield n number of elements from a given list with a balanced distribution
     https://stackoverflow.com/questions/24483182/python-split-list-into-n-chunks#answer-54802737
@@ -203,7 +204,7 @@ class InitElementWorker(BaseWorker):
 
         return list(unique_elements.values())
 
-    def check_worker_activity(self):
+    def check_worker_activity(self) -> bool:
         # Check if workers activity associated to this process is in a pending state
         process = self.request("RetrieveProcess", id=self.process["id"])
         if process.get("activity_state") == ERROR_STATE:
@@ -229,7 +230,7 @@ class InitElementWorker(BaseWorker):
                 raise Exception("Worker activity timeout")
             sleep(timer)
 
-    def run(self):
+    def run(self) -> None:
         elements = self.list_process_elements()
         self.dump_chunks(elements)