From d79ebffe43aa0ffbeec40084f0491fbd35106d28 Mon Sep 17 00:00:00 2001 From: Jusong Yu Date: Fri, 17 Dec 2021 13:13:18 +0100 Subject: [PATCH 1/2] Add double quotes as an optional to bash script escape The escape method will escape the string in the quotes. The defalut implementation is using single quotes to escape the string. However, single quotes will not eval the spacial shell parameters such as $ and @, this makes bash script not versatile to set the parameters in runtime. The `use_double_quotes` option is added to escape the string in the double quotes so the spacial parameters are still valid in the shell script. Using double escape option into environment variables of job templete setting address #4836. User can now use double quotes escaping in calcjob option to escape the environment variables of scheduler. --- aiida/common/escaping.py | 17 ++++++++++--- aiida/engine/processes/calcjobs/calcjob.py | 4 +++ aiida/schedulers/datastructures.py | 4 +++ aiida/schedulers/scheduler.py | 2 +- tests/common/test_escaping.py | 29 +++++++++++++--------- tests/schedulers/test_sge.py | 5 ++-- 6 files changed, 42 insertions(+), 19 deletions(-) diff --git a/aiida/common/escaping.py b/aiida/common/escaping.py index dbec8ba545..170cf5b82c 100644 --- a/aiida/common/escaping.py +++ b/aiida/common/escaping.py @@ -12,7 +12,7 @@ import re -def escape_for_bash(str_to_escape): +def escape_for_bash(str_to_escape, use_double_quotes=False): """ This function takes any string and escapes it in a way that bash will interpret it as a single string. @@ -34,14 +34,23 @@ def escape_for_bash(str_to_escape): Finally, note that for python I have to enclose the string '"'"' within triple quotes to make it work, getting finally: the complicated string found below. + + :param str_to_escape: the string to escape. + :param use_double_quotes: boolean, if ``True``, use double quotes instead of single quotes. + :return: the escaped string. """ if str_to_escape is None: return '' str_to_escape = str(str_to_escape) - - escaped_quotes = str_to_escape.replace("'", """'"'"'""") - return f"'{escaped_quotes}'" + if use_double_quotes: + escaped_quotes = str_to_escape.replace('"', '''"'"'"''') + escaped = f'"{escaped_quotes}"' + else: + escaped_quotes = str_to_escape.replace("'", """'"'"'""") + escaped = f"'{escaped_quotes}'" + + return escaped # Mapping of "SQL" tokens into corresponding regex expressions diff --git a/aiida/engine/processes/calcjobs/calcjob.py b/aiida/engine/processes/calcjobs/calcjob.py index 5231e26e04..86c31ae3b8 100644 --- a/aiida/engine/processes/calcjobs/calcjob.py +++ b/aiida/engine/processes/calcjobs/calcjob.py @@ -244,6 +244,9 @@ def define(cls, spec: CalcJobProcessSpec) -> None: # type: ignore[override] help='If set to true, the submission script will load the system environment variables',) spec.input('metadata.options.environment_variables', valid_type=dict, default=lambda: {}, help='Set a dictionary of custom environment variables for this calculation',) + spec.input('metadata.options.environment_variables_double_quotes', valid_type=bool, default=False, + help='If set to True, use double quotes instead of single quotes to escape the environment variables ' + 'specified in ``environment_variables``.',) spec.input('metadata.options.priority', valid_type=str, required=False, help='Set the priority of the job to be queued') spec.input('metadata.options.max_memory_kb', valid_type=int, required=False, @@ -731,6 +734,7 @@ def presubmit(self, folder: Folder) -> CalcInfo: job_tmpl.import_sys_environment = self.node.get_option('import_sys_environment') job_tmpl.job_environment = self.node.get_option('environment_variables') + job_tmpl.environment_variables_double_quotes = self.node.get_option('environment_variables_double_quotes') queue_name = self.node.get_option('queue_name') account = self.node.get_option('account') diff --git a/aiida/schedulers/datastructures.py b/aiida/schedulers/datastructures.py index c44ad70784..2ef978991d 100644 --- a/aiida/schedulers/datastructures.py +++ b/aiida/schedulers/datastructures.py @@ -248,6 +248,9 @@ class JobTemplate(DefaultFieldsAttributeDict): # pylint: disable=too-many-insta * ``rerunnable``: if the job is rerunnable (boolean) * ``job_environment``: a dictionary with environment variables to set before the execution of the code. + * ``environment_variables_double_quotes``: if set to True, use double quotes + instead of single quotes to escape the environment variables specified + in ``environment_variables``. * ``working_directory``: the working directory for this job. During submission, the transport will first do a 'chdir' to this directory, and then possibly set a scheduler parameter, if this is supported @@ -328,6 +331,7 @@ class JobTemplate(DefaultFieldsAttributeDict): # pylint: disable=too-many-insta 'submit_as_hold', 'rerunnable', 'job_environment', + 'environment_variables_double_quotes', 'working_directory', 'email', 'email_on_started', diff --git a/aiida/schedulers/scheduler.py b/aiida/schedulers/scheduler.py index 867964870a..35f258070f 100644 --- a/aiida/schedulers/scheduler.py +++ b/aiida/schedulers/scheduler.py @@ -182,7 +182,7 @@ def _get_submit_script_environment_variables(self, template): # pylint: disable lines = ['# ENVIRONMENT VARIABLES BEGIN ###'] for key, value in template.job_environment.items(): - lines.append(f'export {key.strip()}={escape_for_bash(value)}') + lines.append(f'export {key.strip()}={escape_for_bash(value, template.environment_variables_double_quotes)}') lines.append('# ENVIRONMENT VARIABLES END ###') diff --git a/tests/common/test_escaping.py b/tests/common/test_escaping.py index 80bdb377c6..6d6b7c45e3 100644 --- a/tests/common/test_escaping.py +++ b/tests/common/test_escaping.py @@ -8,19 +8,24 @@ # For further information please visit http://www.aiida.net # ########################################################################### """Tests for the :mod:`aiida.common.escaping`.""" +import pytest + from aiida.common.escaping import escape_for_bash -def test_escape_for_bash(): +@pytest.mark.parametrize(('to_escape, expected_single_quotes, expected_double_quotes'), ( + (None, '', ''), + ('string', "'string'", '"string"'), + ('string with space', "'string with space'", '"string with space"'), + ( + """string with ' single and " double quote""", """'string with '"'"' single and " double quote'""", + '''"string with ' single and "'"'" double quote"''' + ), + (1, "'1'", '"1"'), + (2.0, "'2.0'", '"2.0"'), + ('$PWD', "'$PWD'", '"$PWD"'), +)) +def test_escape_for_bash(to_escape, expected_single_quotes, expected_double_quotes): """Tests various inputs for `aiida.common.escaping.escape_for_bash`.""" - tests = ( - [None, ''], - ['string', "'string'"], - ['string with space', "'string with space'"], - ["string with a ' single quote", """'string with a '"'"' single quote'"""], - [1, "'1'"], - [2.0, "'2.0'"], - ) - - for string_input, string_escaped in tests: - assert escape_for_bash(string_input) == string_escaped + assert escape_for_bash(to_escape, use_double_quotes=False) == expected_single_quotes + assert escape_for_bash(to_escape, use_double_quotes=True) == expected_double_quotes diff --git a/tests/schedulers/test_sge.py b/tests/schedulers/test_sge.py index d92e2b2cab..67ec7230d3 100644 --- a/tests/schedulers/test_sge.py +++ b/tests/schedulers/test_sge.py @@ -327,6 +327,7 @@ def test_submit_script(self): job_tmpl.priority = None job_tmpl.max_wallclock_seconds = '3600' # "23:59:59" job_tmpl.job_environment = {'HOME': '/home/users/dorigm7s/', 'WIENROOT': '$HOME:/WIEN2k'} + job_tmpl.environment_variables_double_quotes = True submit_script_text = sge.get_submit_script(job_tmpl) @@ -335,8 +336,8 @@ def test_submit_script(self): self.assertTrue('#$ -q FavQ.q' in submit_script_text) self.assertTrue('#$ -l h_rt=01:00:00' in submit_script_text) self.assertTrue('# ENVIRONMENT VARIABLES BEGIN ###' in submit_script_text) - self.assertTrue("export HOME='/home/users/dorigm7s/'" in submit_script_text) - self.assertTrue("export WIENROOT='$HOME:/WIEN2k'" in submit_script_text) + self.assertTrue('export HOME="/home/users/dorigm7s/"' in submit_script_text) + self.assertTrue('export WIENROOT="$HOME:/WIEN2k"' in submit_script_text) self.assertFalse('#$ -r yes' in submit_script_text) def test_submit_script_rerunnable(self): # pylint: disable=no-self-use From a95c09a3325738ceb0af55becddf65ea475b38ef Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 13 Feb 2022 14:35:50 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- aiida/schedulers/datastructures.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiida/schedulers/datastructures.py b/aiida/schedulers/datastructures.py index 2ef978991d..2ca464f925 100644 --- a/aiida/schedulers/datastructures.py +++ b/aiida/schedulers/datastructures.py @@ -248,8 +248,8 @@ class JobTemplate(DefaultFieldsAttributeDict): # pylint: disable=too-many-insta * ``rerunnable``: if the job is rerunnable (boolean) * ``job_environment``: a dictionary with environment variables to set before the execution of the code. - * ``environment_variables_double_quotes``: if set to True, use double quotes - instead of single quotes to escape the environment variables specified + * ``environment_variables_double_quotes``: if set to True, use double quotes + instead of single quotes to escape the environment variables specified in ``environment_variables``. * ``working_directory``: the working directory for this job. During submission, the transport will first do a 'chdir' to this directory,