allow int env values in job spec

Author: andre-merzky

src/psij/executors/batch/batch_scheduler_executor.py

@@ -53,7 +53,8 @@ def _env_to_mustache(job: Job) -> List[Dict[str, str]]:
 
     r = []
     for k, v in job.spec.environment.items():
-        r.append({'name': k, 'value': bash_escape(v)})
+        r.append({'name': k, 'value': bash_escape(str(v))})
+
     return r
 
 

src/psij/executors/local.py

@@ -98,6 +98,10 @@ def _get_env(spec: JobSpec) -> Optional[Dict[str, str]]:
             # merge current env with spec env
             env = os.environ.copy()
             env.update(spec.environment)
+            # convert integers to strings
+            for k, v in env.items():
+                if isinstance(v, int):
+                    env[k] = str(v)
             return env
     else:
         # only spec env

src/psij/job_executor.py

@@ -67,7 +67,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
@@ -101,9 +101,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__))
 
         job.executor = self
         return spec

src/psij/job_spec.py

@@ -1,6 +1,6 @@
 import sys
 from pathlib import Path
-from typing import Optional, List, Dict, Any
+from typing import Optional, List, Dict, Any, Union
 
 from typeguard import check_argument_types
 
@@ -15,7 +15,8 @@ class JobSpec(object):
 
     def __init__(self, name: Optional[str] = None, executable: Optional[str] = None,
                  arguments: Optional[List[str]] = None, directory: Optional[Path] = 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: Optional[Path] = None, stdout_path: Optional[Path] = None,
                  stderr_path: Optional[Path] = None, resources: Optional[ResourceSpec] = None,
                  attributes: Optional[JobAttributes] = None, pre_launch: Optional[Path] = None,

tests/plugins1/_batch_test/test/test.mustache

@@ -22,8 +22,8 @@ export {{key}}="{{value}}"
     {{/custom_attributes.test}}
 {{/job.spec.attributes}}
 
-{{#job.spec.inherit_environment}}env \{{/job.spec.inherit_environment}}{{^job.spec.inherit_environment}}env --ignore-environment \{{/job.spec.inherit_environment}}{{#env}}
-{{name}}="{{value}}" \
+{{#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}}
 
-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

@@ -112,7 +112,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

@@ -117,18 +117,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_'}
         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

@@ -10,16 +10,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()
@@ -32,8 +30,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}
     assert spec.environment['foo'] == 'biz'
+    assert spec.environment['bar'] == 42
 
 
 test_environment_types()