Use IO[str] rather than TextIO in batch scheduler executors in the context

of submit script generation.

Author: hategan

src/psij/executors/batch/batch_scheduler_executor.py

@@ -7,7 +7,7 @@
 from datetime import timedelta
 from pathlib import Path
 from threading import Thread, RLock
-from typing import Optional, List, Dict, Collection, cast, TextIO, Union
+from typing import Optional, List, Dict, Collection, cast, Union, IO
 
 from .escape_functions import bash_escape
 from psij.launchers.script_based_launcher import ScriptBasedLauncher
@@ -268,7 +268,7 @@ def attach(self, job: Job, native_id: str) -> None:
 
     @abstractmethod
     def generate_submit_script(self, job: Job, context: Dict[str, object],
-                               submit_file: TextIO) -> None:
+                               submit_file: IO[str]) -> None:
         """Called to generate a submit script for a job.
 
         Concrete implementations of batch scheduler executors must override this method in

src/psij/executors/batch/cobalt.py

@@ -1,7 +1,7 @@
 """Defines a JobExecutor for the Cobalt resource manager."""
 from datetime import timedelta
 from pathlib import Path
-from typing import Optional, Collection, List, Dict, TextIO
+from typing import Optional, Collection, List, Dict, IO
 import re
 import os
 import stat
@@ -59,7 +59,7 @@ def __init__(
         )
 
     def generate_submit_script(
-        self, job: Job, context: Dict[str, object], submit_file: TextIO
+        self, job: Job, context: Dict[str, object], submit_file: IO[str]
     ) -> None:
         """See :meth:`~.BatchSchedulerExecutor.generate_submit_script`."""
         self.generator.generate_submit_script(job, context, submit_file)

src/psij/executors/batch/lsf.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 import re
 import json
-from typing import Optional, Collection, List, Dict, TextIO
+from typing import Optional, Collection, List, Dict, IO
 
 from psij import Job, JobStatus, JobState, SubmitException
 from psij.executors.batch.batch_scheduler_executor import (
@@ -62,7 +62,7 @@ def __init__(self, url: Optional[str], config: Optional[LsfExecutorConfig] = Non
         )
 
     def generate_submit_script(
-        self, job: Job, context: Dict[str, object], submit_file: TextIO
+        self, job: Job, context: Dict[str, object], submit_file: IO[str]
     ) -> None:
         """See :meth:`~.BatchSchedulerExecutor.generate_submit_script`."""
         assert job.spec is not None

src/psij/executors/batch/pbspro.py

@@ -1,6 +1,6 @@
 from datetime import timedelta
 from pathlib import Path
-from typing import Optional, Collection, List, Dict, TextIO
+from typing import Optional, Collection, List, Dict, IO
 
 from psij import Job, JobStatus, JobState, SubmitException
 from psij.executors.batch.batch_scheduler_executor import BatchSchedulerExecutor, \
@@ -72,7 +72,7 @@ def __init__(self, url: Optional[str] = None, config: Optional[PBSProExecutorCon
     # Submit methods
 
     def generate_submit_script(self, job: Job, context: Dict[str, object],
-                               submit_file: TextIO) -> None:
+                               submit_file: IO[str]) -> None:
         """See :meth:`~.BatchSchedulerExecutor.generate_submit_script`."""
         self.generator.generate_submit_script(job, context, submit_file)
 

src/psij/executors/batch/script_generator.py

@@ -1,6 +1,6 @@
 import pathlib
 from abc import ABC
-from typing import TextIO, Dict, Callable
+from typing import Dict, Callable, IO
 
 import pystache
 
@@ -29,7 +29,7 @@ def __init__(self, config: JobExecutorConfig) -> None:
         """
         self.config = config
 
-    def generate_submit_script(self, job: Job, context: Dict[str, object], out: TextIO) -> None:
+    def generate_submit_script(self, job: Job, context: Dict[str, object], out: IO[str]) -> None:
         """Generates a job submit script.
 
         Concerete implementations of submit script generators must implement this method. Its
@@ -76,7 +76,7 @@ def __init__(self, config: JobExecutorConfig, template_path: pathlib.Path,
             self.template = pystache.parse(template_file.read())
         self.renderer = pystache.Renderer(escape=escape)
 
-    def generate_submit_script(self, job: Job, context: Dict[str, object], out: TextIO) -> None:
+    def generate_submit_script(self, job: Job, context: Dict[str, object], out: IO[str]) -> None:
         """See :func:`~SubmitScriptGenerator.generate_submit_script`.
 
         Renders a submit script using the template specified when this generator was constructed.

src/psij/executors/batch/slurm.py

@@ -1,6 +1,6 @@
 from datetime import timedelta
 from pathlib import Path
-from typing import Optional, Collection, List, Dict, TextIO
+from typing import Optional, Collection, List, Dict, IO
 
 from psij import Job, JobStatus, JobState, SubmitException
 from psij.executors.batch.batch_scheduler_executor import BatchSchedulerExecutor, \
@@ -117,7 +117,7 @@ def __init__(self, url: Optional[str] = None, config: Optional[SlurmExecutorConf
                                                   / 'slurm.mustache')
 
     def generate_submit_script(self, job: Job, context: Dict[str, object],
-                               submit_file: TextIO) -> None:
+                               submit_file: IO[str]) -> None:
         """See :meth:`~.BatchSchedulerExecutor.generate_submit_script`."""
         self.generator.generate_submit_script(job, context, submit_file)
 

tests/plugins1/_batch_test/_batch_test.py

@@ -1,6 +1,6 @@
 import sys
 from pathlib import Path
-from typing import Optional, Collection, List, Dict, TextIO, cast
+from typing import Optional, Collection, List, Dict, IO, cast
 
 from psij import Job, JobStatus, JobState, SubmitException, JobExecutorConfig, ResourceSpecV1
 from psij.executors.batch.batch_scheduler_executor import BatchSchedulerExecutor, \
@@ -48,7 +48,7 @@ def __init__(self, url: Optional[str] = None, config: Optional[_TestExecutorConf
                                                   / 'test.mustache')
 
     def generate_submit_script(self, job: Job, context: Dict[str, object],
-                               submit_file: TextIO) -> None:
+                               submit_file: IO[str]) -> None:
         self.generator.generate_submit_script(job, context, submit_file)
 
     def get_submit_command(self, job: Job, submit_file_path: Path) -> List[str]: