Repository: DataBiosphere/toil

Issue:
Add type hints to worker.py
Add type hints to `src/toil/worker.py` so it can be checked under mypy during linting.

Refers to #3568.

┆Issue is synchronized with this [Jira Task](https://ucsc-cgl.atlassian.net/browse/TOIL-947)
┆Issue Number: TOIL-947



Discussion:

Resolving patch:
diff --git a/contrib/admin/mypy-with-ignore.py b/contrib/admin/mypy-with-ignore.py
--- a/contrib/admin/mypy-with-ignore.py
+++ b/contrib/admin/mypy-with-ignore.py
@@ -29,7 +29,6 @@ def main():
         'src/toil/job.py',
         'src/toil/leader.py',
         'src/toil/common.py',
-        'src/toil/worker.py',
         'src/toil/__init__.py',
         'src/toil/resource.py',
         'src/toil/deferred.py',
diff --git a/src/toil/jobStores/abstractJobStore.py b/src/toil/jobStores/abstractJobStore.py
--- a/src/toil/jobStores/abstractJobStore.py
+++ b/src/toil/jobStores/abstractJobStore.py
@@ -1078,7 +1078,8 @@ def updateFile(self, jobStoreFileID: str, localFilePath: str) -> None:
         raise NotImplementedError()
 
     @abstractmethod
-    def updateFileStream(self, jobStoreFileID: str, encoding: Optional[str] = None, errors: Optional[str] = None) -> None:
+    @contextmanager
+    def updateFileStream(self, jobStoreFileID: str, encoding: Optional[str] = None, errors: Optional[str] = None) -> Iterator[IO[Any]]:
         """
         Replaces the existing version of a file in the job store. Similar to writeFile, but
         returns a context manager yielding a file handle which can be written to. The
diff --git a/src/toil/worker.py b/src/toil/worker.py
--- a/src/toil/worker.py
+++ b/src/toil/worker.py
@@ -25,13 +25,16 @@
 import sys
 import time
 import traceback
+
 from contextlib import contextmanager
+from typing import List, Iterator, Optional
 
 from toil import logProcessContext
-from toil.common import Toil, safeUnpickleFromStream
+from toil.common import Toil, Config, safeUnpickleFromStream
 from toil.deferred import DeferredFunctionManager
 from toil.fileStores.abstractFileStore import AbstractFileStore
-from toil.job import CheckpointJobDescription, Job
+from toil.job import CheckpointJobDescription, Job, JobDescription
+from toil.jobStores.abstractJobStore import AbstractJobStore
 from toil.lib.expando import MagicExpando
 from toil.lib.io import make_public_dir
 from toil.lib.resources import (get_total_cpu_time,
@@ -47,15 +50,14 @@
 logger = logging.getLogger(__name__)
 
 
-def nextChainable(predecessor, jobStore, config):
+def nextChainable(predecessor: JobDescription, jobStore: AbstractJobStore, config: Config) -> Optional[JobDescription]:
     """
     Returns the next chainable job's JobDescription after the given predecessor
     JobDescription, if one exists, or None if the chain must terminate.
 
-    :param toil.job.JobDescription predecessor: The job to chain from
-    :param toil.jobStores.abstractJobStore.AbstractJobStore jobStore: The JobStore to fetch JobDescriptions from.
-    :param toil.common.Config config: The configuration for the current run.
-    :rtype: toil.job.JobDescription or None
+    :param predecessor: The job to chain from
+    :param jobStore: The JobStore to fetch JobDescriptions from.
+    :param config: The configuration for the current run.
     """
     #If no more jobs to run or services not finished, quit
     if len(predecessor.stack) == 0 or len(predecessor.services) > 0 or (isinstance(predecessor, CheckpointJobDescription) and predecessor.checkpoint != None):
@@ -121,13 +123,14 @@ def nextChainable(predecessor, jobStore, config):
     # Made it through! This job is chainable.
     return successor
 
-def workerScript(jobStore, config, jobName, jobStoreID, redirectOutputToLogFile=True):
+def workerScript(jobStore: AbstractJobStore, config: Config, jobName: str, jobStoreID: str, redirectOutputToLogFile: bool = True) -> int:
     """
     Worker process script, runs a job.
 
-    :param str jobName: The "job name" (a user friendly name) of the job to be run
-    :param str jobStoreLocator: Specifies the job store to use
-    :param str jobStoreID: The job store ID of the job to be run
+    :param jobStore: The JobStore to fetch JobDescriptions from.
+    :param config: The configuration for the current run.
+    :param jobName: The "job name" (a user friendly name) of the job to be run
+    :param jobStoreID: The job store ID of the job to be run
 
     :return int: 1 if a job failed, or 0 if all jobs succeeded
     """
@@ -582,7 +585,7 @@ def workerScript(jobStore, config, jobName, jobStoreID, redirectOutputToLogFile=
         statsDict.logs.messages = logMessages
 
     if (debugging or config.stats or statsDict.workers.logsToMaster) and not jobAttemptFailed:  # We have stats/logging to report back
-        jobStore.writeStatsAndLogging(json.dumps(statsDict, ensure_ascii=True).encode())
+        jobStore.writeStatsAndLogging(json.dumps(statsDict, ensure_ascii=True))
 
     #Remove the temp dir
     cleanUp = config.cleanWorkDir
@@ -601,7 +604,7 @@ def workerScript(jobStore, config, jobName, jobStoreID, redirectOutputToLogFile=
     else:
         return 0
 
-def parse_args(args):
+def parse_args(args: List[str]) -> argparse.Namespace:
     """
     Parse command-line arguments to the worker.
     """
@@ -634,7 +637,7 @@ def parse_args(args):
 
 
 @contextmanager
-def in_contexts(contexts):
+def in_contexts(contexts: List[str]) -> Iterator[None]:
     """
     Unpickle and enter all the pickled, base64-encoded context managers in the
     given list. Then do the body, then leave them all.
@@ -661,7 +664,7 @@ def in_contexts(contexts):
                 yield
 
 
-def main(argv=None):
+def main(argv: Optional[List[str]] = None) -> None:
     if argv is None:
         argv = sys.argv
 


Test patch: