Merge pull request #425 from ExaWorks/fix-env-type

Fix env type

Author: andre-merzky

src/psij/job_executor.py

@@ -65,7 +65,7 @@ def _check_job(self, job: Job) -> JobSpec:
         Verifies that various aspects of the job are correctly specified. This includes precisely
         the following checks:
         * the job has a non-null specification
-        * job.spec.environment is a Dict[str, str]
+        * job.spec.environment is a Dict[str, [str | int]]
 
         While this method makes a fair attempt at ensuring the validity of the job, it makes no
         such guarantees. Specifically, if an executor implementation requires checks not listed
@@ -99,9 +99,9 @@ def _check_job(self, job: Job) -> JobSpec:
                     if not isinstance(k, str):
                         raise TypeError('environment key "%s" is not a string (%s)'
                                         % (k, type(k).__name__))
-                    if not isinstance(v, str):
-                        raise TypeError('environment key "%s" has non-string value (%s)'
-                                        % (k, type(v).__name__))
+                    if not isinstance(v, (str, int)):
+                        raise TypeError('environment value for key "%s" must be string '
+                                        'or int type (%s)' % (k, type(v).__name__))
 
         if job.executor is not None:
             raise InvalidJobException('Job is already associated with an executor')

src/psij/job_spec.py

@@ -21,6 +21,18 @@ def _to_path(arg: Union[str, pathlib.Path, None]) -> Optional[pathlib.Path]:
         return pathlib.Path(arg)
 
 
+def _to_env_dict(arg: Union[Dict[str, Union[str, int]], None]) -> Optional[Dict[str, str]]:
+    if arg is None:
+        return None
+    ret = dict()
+    for k, v in arg.items():
+        if isinstance(v, int):
+            ret[k] = str(v)
+        else:
+            ret[k] = v
+    return ret
+
+
 class JobSpec(object):
     """A class that describes the details of a job."""
 
@@ -29,7 +41,8 @@ def __init__(self, executable: Optional[str] = None, arguments: Optional[List[st
                  # sphinx fails to find the class. Using Path in the getters and setters does not
                  # appear to trigger a problem.
                  directory: Union[str, pathlib.Path, None] = None, name: Optional[str] = None,
-                 inherit_environment: bool = True, environment: Optional[Dict[str, str]] = None,
+                 inherit_environment: bool = True,
+                 environment: Optional[Dict[str, Union[str, int]]] = None,
                  stdin_path: Union[str, pathlib.Path, None] = None,
                  stdout_path: Union[str, pathlib.Path, None] = None,
                  stderr_path: Union[str, pathlib.Path, None] = None,
@@ -125,7 +138,7 @@ def __init__(self, executable: Optional[str] = None, arguments: Optional[List[st
         # care of the conversion, but mypy gets confused
         self._directory = _to_path(directory)
         self.inherit_environment = inherit_environment
-        self.environment = environment
+        self.environment = _to_env_dict(environment)
         self._stdin_path = _to_path(stdin_path)
         self._stdout_path = _to_path(stdout_path)
         self._stderr_path = _to_path(stderr_path)
@@ -152,6 +165,16 @@ def name(self) -> Optional[str]:
     def name(self, value: Optional[str]) -> None:
         self._name = value
 
+    @property
+    def environment(self) -> Optional[Dict[str, str]]:
+        """Return the environment dict."""
+        return self._environment
+
+    @environment.setter
+    def environment(self, env: Optional[Dict[str, Union[str, int]]]) -> None:
+        """Ensure env dict values to be string typed."""
+        self._environment = _to_env_dict(env)
+
     @property
     def directory(self) -> Optional[pathlib.Path]:
         """The directory, on the compute side, in which the executable is to be run."""

tests/plugins1/_batch_test/test/test.mustache

@@ -43,7 +43,7 @@ done
 export PSIJ_NODEFILE
 
 {{#job.spec.inherit_environment}}env \{{/job.spec.inherit_environment}}{{^job.spec.inherit_environment}}env --ignore-environment \{{/job.spec.inherit_environment}}{{#env}}
-{{name}}="{{value}}" \
-{{/env}}{{#psij.launch_command}}{{.}} {{/psij.launch_command}}
+{{name}}="{{value}}" \{{/env}}
+{{#psij.launch_command}}{{.}} {{/psij.launch_command}}
 
-echo "$?" > "{{psij.script_dir}}/$PSIJ_BATCH_TEST_JOB_ID.ec"
+echo "$?" > "{{psij.script_dir}}/$PSIJ_BATCH_TEST_JOB_ID.ec"

tests/test_doc_examples.py

@@ -114,7 +114,7 @@ def test_job_parameters() -> None:
             psij.JobSpec(
                 executable="/bin/hostname",
                 stdout_path=output_path,
-                environment={"FOOBAR": "BAZ"},  # custom environment has no effect here
+                environment={"FOOBAR": "BAZ", "BUZ": 1},  # custom environment has no effect here
                 directory=pathlib.Path(td),  # CWD has no effect on result here
             )
         )

tests/test_executor.py

@@ -122,18 +122,19 @@ def test_env_var(execparams: ExecutorTestParams) -> None:
     _make_test_dir()
     with TemporaryDirectory(dir=Path.home() / '.psij' / 'test') as td:
         outp = Path(td, 'stdout.txt')
-        job = Job(JobSpec(executable='/bin/bash', arguments=['-c', 'echo -n $TEST_VAR'],
+        job = Job(JobSpec(executable='/bin/bash',
+                          arguments=['-c', 'env > /tmp/t; echo -n $TEST_VAR$TEST_INT'],
                           stdout_path=outp))
         assert job.spec is not None
-        job.spec.environment = {'TEST_VAR': '_y_'}
+        job.spec.environment = {'TEST_INT': 1, 'TEST_VAR': '_y_'}  # type: ignore
         ex = _get_executor_instance(execparams, job)
         ex.submit(job)
         status = job.wait(timeout=_get_timeout(execparams))
         assert_completed(job, status)
         f = outp.open("r")
         contents = f.read()
         f.close()
-        assert contents == '_y_'
+        assert contents == '_y_1'
 
 
 def test_stdin_redirect(execparams: ExecutorTestParams) -> None:

tests/test_job_spec.py

@@ -12,16 +12,14 @@ def _test_spec(spec: JobSpec) -> None:
 
 
 def test_environment_types() -> None:
-    with pytest.raises(TypeError):
-        _test_spec(JobSpec(environment={'foo': 1}))  # type: ignore
 
     with pytest.raises(TypeError):
-        _test_spec(JobSpec(environment={1: 'foo'}))  # type: ignore
+        _test_spec(JobSpec(executable='true', environment={1: 'foo'}))  # type: ignore
 
     with pytest.raises(TypeError):
-        spec = JobSpec()
+        spec = JobSpec(executable='true')
         spec.environment = {'foo': 'bar'}
-        spec.environment['buz'] = 2  # type: ignore
+        spec.environment['buz'] = [2]  # type: ignore
         _test_spec(spec)
 
     spec = JobSpec()
@@ -34,8 +32,9 @@ def test_environment_types() -> None:
     spec.environment = {'foo': 'bar'}
     assert spec.environment['foo'] == 'bar'
 
-    spec.environment = {'foo': 'biz'}
+    spec.environment = {'foo': 'biz', 'bar': 42}  # type: ignore
     assert spec.environment['foo'] == 'biz'
+    assert spec.environment['bar'] == '42'
 
 
 def test_path_conversion() -> None: