diff --git a/src/psij/job_executor.py b/src/psij/job_executor.py index f763e2da..b9f9a226 100644 --- a/src/psij/job_executor.py +++ b/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') diff --git a/src/psij/job_spec.py b/src/psij/job_spec.py index 25cdf44e..7eec03ec 100644 --- a/src/psij/job_spec.py +++ b/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.""" diff --git a/tests/plugins1/_batch_test/test/test.mustache b/tests/plugins1/_batch_test/test/test.mustache index 245e3ba8..ae51ae68 100644 --- a/tests/plugins1/_batch_test/test/test.mustache +++ b/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" \ No newline at end of file +echo "$?" > "{{psij.script_dir}}/$PSIJ_BATCH_TEST_JOB_ID.ec" diff --git a/tests/test_doc_examples.py b/tests/test_doc_examples.py index 752af818..87e95c23 100644 --- a/tests/test_doc_examples.py +++ b/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 ) ) diff --git a/tests/test_executor.py b/tests/test_executor.py index 137d0780..e425a030 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -122,10 +122,11 @@ 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)) @@ -133,7 +134,7 @@ def test_env_var(execparams: ExecutorTestParams) -> None: f = outp.open("r") contents = f.read() f.close() - assert contents == '_y_' + assert contents == '_y_1' def test_stdin_redirect(execparams: ExecutorTestParams) -> None: diff --git a/tests/test_job_spec.py b/tests/test_job_spec.py index b6a61a3e..9fa57a49 100644 --- a/tests/test_job_spec.py +++ b/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: