From 58694e6f71b4efa4bad47a5bfbd2928d3c209d12 Mon Sep 17 00:00:00 2001
From: Albert Zeyer <albert.zeyer@rwth-aachen.de>
Date: Wed, 3 Jul 2024 13:49:59 +0200
Subject: [PATCH] typing updates for Task (#194)

Also removed Python 2 style code to derive from object.

I did not cover update_rqmt yet
because this would require Protocol (not just Callable),
and I'm not sure if we want that (requires Python 3.8).
---
 sisyphus/task.py | 40 ++++++++++++++++++++--------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/sisyphus/task.py b/sisyphus/task.py
index 0218fc7..b579dd0 100644
--- a/sisyphus/task.py
+++ b/sisyphus/task.py
@@ -3,7 +3,7 @@
 import logging
 import sys
 import time
-import typing
+from typing import Optional, Union, Any, Sequence, Dict, List
 import subprocess as sp
 from ast import literal_eval
 
@@ -11,27 +11,27 @@
 import sisyphus.global_settings as gs
 
 
-class Task(object):
+class Task:
     """
     Object to hold information what function should be run with which requirements.
     """
 
     def __init__(
         self,
-        start,
-        resume=None,
-        rqmt=None,
-        args=None,
-        mini_task=False,
+        start: str,
+        resume: Optional[str] = None,
+        rqmt: Optional[Dict[str, Any]] = None,
+        args: Optional[Sequence[Union[List[Any], Any]]] = None,
+        mini_task: bool = False,
         update_rqmt=None,
-        parallel=0,
-        tries=1,
-        continuable=False,
+        parallel: int = 0,
+        tries: int = 1,
+        continuable: bool = False,
     ):
         """
-        :param str start: name of the function which will be executed on start
-        :param str resume: name of the function which will be executed on resume, often set equal to start
-        :param dict[str] rqmt: job requirements
+        :param start: name of the function which will be executed on start
+        :param resume: name of the function which will be executed on resume, often set equal to start
+        :param rqmt: job requirements
             Might contain:
                 "cpu": number of cpus
                 "gpu": number of gpus
@@ -40,14 +40,14 @@ def __init__(
                 "multi_node_slots": amount of slots, distributed potentially over multiple nodes.
                                     E.g. maps to `--ntasks <multi_node_slots>` parameter for Slurm,
                                     and to `-pe <pe_name> <multi_node_slots>` for SGE (SGE parallel environment (PE)).
-        :param typing.Sequence[typing.Union[typing.List[object],object]] args: job arguments
-        :param bool mini_task: will be run on engine for short jobs if True
+        :param args: job arguments
+        :param mini_task: will be run on engine for short jobs if True
         :param (dict[str],dict[str])->dict[str] update_rqmt: function to update job requirements for interrupted jobs
-        :param int parallel: if set to > 0, groups jobs for individual arguments together into the number of batches
-                             specified here. Will then submit at max `parallel` jobs into the engine at a time.
-        :param int tries: how often this task is resubmitted after failure
-        :param bool continuable: If set to True this task will not set a finished marker, useful for tasks that can be
-                                 continued for arbitrarily long, e.g. adding more epochs to neural network training
+        :param parallel: if set to > 0, groups jobs for individual arguments together into the number of batches
+                         specified here. Will then submit at max `parallel` jobs into the engine at a time.
+        :param tries: how often this task is resubmitted after failure
+        :param continuable: If set to True this task will not set a finished marker, useful for tasks that can be
+                            continued for arbitrarily long, e.g. adding more epochs to neural network training
         """
         if rqmt is None:
             rqmt = {}