From 42203bab35d254773c6099c35c30b2f167a71904 Mon Sep 17 00:00:00 2001 From: yacchin1205 <968739+yacchin1205@users.noreply.github.com> Date: Mon, 3 Mar 2025 17:03:18 +0900 Subject: [PATCH 1/2] Add function to suppress parameter output of sensitive information --- docs/reference/papermill-cli.rst | 6 + docs/usage-cli.rst | 6 + papermill/cli.py | 22 ++++ papermill/engines.py | 8 ++ papermill/execute.py | 11 +- papermill/parameterize.py | 22 +++- papermill/tests/test_cli.py | 160 +++++++++++++++++++++++++++ papermill/tests/test_parameterize.py | 23 ++++ papermill/tests/test_utils.py | 30 +++++ papermill/utils.py | 76 +++++++++++++ 10 files changed, 361 insertions(+), 3 deletions(-) diff --git a/docs/reference/papermill-cli.rst b/docs/reference/papermill-cli.rst index 09ba7eb2..7beb79ca 100644 --- a/docs/reference/papermill-cli.rst +++ b/docs/reference/papermill-cli.rst @@ -71,5 +71,11 @@ Command Line options failing execution (default: forever) --report-mode / --no-report-mode Flag for hiding input. + --obfuscate-sensitive-parameters / --no-obfuscate-sensitive-parameters + Flag for obfuscating sensitive parameters. + --sensitive-parameter-patterns TEXT... + List of patterns for obfuscating parameter + names in notebooks. If not provided, defaults + to papermill.utils.SENSITIVE_PARAMETER_PATTERNS. --version Flag for displaying the version. -h, --help Show this message and exit. diff --git a/docs/usage-cli.rst b/docs/usage-cli.rst index f0088488..5579e317 100644 --- a/docs/usage-cli.rst +++ b/docs/usage-cli.rst @@ -78,5 +78,11 @@ options: --report-mode / --no-report-mode Flag for hiding input. + --obfuscate-sensitive-parameters / --no-obfuscate-sensitive-parameters + Flag for obfuscating sensitive parameters. + --sensitive-parameter-patterns TEXT... + List of patterns for obfuscating parameter + names in notebooks. If not provided, defaults + to papermill.utils.SENSITIVE_PARAMETER_PATTERNS. --version Flag for displaying the version. -h, --help Show this message and exit. diff --git a/papermill/cli.py b/papermill/cli.py index 0c1b54d0..aca8d8f9 100755 --- a/papermill/cli.py +++ b/papermill/cli.py @@ -15,6 +15,7 @@ from .execute import execute_notebook from .inspection import display_notebook_help from .iorw import NoDatesSafeLoader, read_yaml_file +from .utils import SENSITIVE_PARAMETER_PATTERNS from .version import version as papermill_version click.disable_unicode_literals_warning = True @@ -129,6 +130,19 @@ def print_papermill_version(ctx, param, value): help="Time in seconds to wait for each cell before failing execution (default: forever)", ) @click.option('--report-mode/--no-report-mode', default=False, help="Flag for hiding input.") +@click.option( + '--obfuscate-sensitive-parameters/--no-obfuscate-sensitive-parameters', + default=True, + help="Flag for obfuscating sensitive parameters.", +) +@click.option( + '--sensitive-parameter-patterns', + multiple=True, + help=( + "List of patterns for obfuscating parameter names in notebooks." + "If not provided, defaults to papermill.utils.SENSITIVE_PARAMETER_PATTERNS." + ), +) @click.option( '--version', is_flag=True, @@ -163,6 +177,8 @@ def papermill( start_timeout, execution_timeout, report_mode, + obfuscate_sensitive_parameters, + sensitive_parameter_patterns, stdout_file, stderr_file, ): @@ -228,6 +244,10 @@ def papermill( for name, value in parameters_raw or []: parameters_final[name] = value + if sensitive_parameter_patterns is not None and len(sensitive_parameter_patterns) == 0: + # If the user does not provide any patterns, we should use the default patterns + sensitive_parameter_patterns = None + if help_notebook: sys.exit(display_notebook_help(click_ctx, notebook_path, parameters_final)) @@ -248,6 +268,8 @@ def papermill( stderr_file=stderr_file, start_timeout=start_timeout, report_mode=report_mode, + obfuscate_sensitive_parameters=obfuscate_sensitive_parameters, + sensitive_parameter_patterns=sensitive_parameter_patterns, cwd=cwd, execution_timeout=execution_timeout, ) diff --git a/papermill/engines.py b/papermill/engines.py index 4d096ba5..08e578ae 100644 --- a/papermill/engines.py +++ b/papermill/engines.py @@ -372,6 +372,14 @@ def execute_notebook( nb_man.cleanup_pbar() nb_man.notebook_complete() + # Replace the source with the obfuscated content if it is in the metadata. + for cell in nb_man.nb.cells: + if cell.get('cell_type') != 'code': + continue + if 'papermill-obfuscated-source' not in cell.metadata: + continue + cell.source = cell.metadata['papermill-obfuscated-source'] + del cell.metadata['papermill-obfuscated-source'] return nb_man.nb @classmethod diff --git a/papermill/execute.py b/papermill/execute.py index eb76cc34..3674daf8 100644 --- a/papermill/execute.py +++ b/papermill/execute.py @@ -8,7 +8,7 @@ from .iorw import get_pretty_path, load_notebook_node, local_file_io_cwd, write_ipynb from .log import logger from .parameterize import add_builtin_parameters, parameterize_notebook, parameterize_path -from .utils import chdir +from .utils import chdir, obfuscate_parameters def execute_notebook( @@ -27,6 +27,8 @@ def execute_notebook( start_timeout=60, report_mode=False, cwd=None, + obfuscate_sensitive_parameters=True, + sensitive_parameter_patterns=None, **engine_kwargs, ): """Executes a single notebook locally. @@ -61,6 +63,11 @@ def execute_notebook( Flag for whether or not to hide input. cwd : str or Path, optional Working directory to use when executing the notebook + obfuscate_sensitive_parameters : bool, optional + Obfuscate sensitive parameters in the notebook, Defaults to True + sensitive_parameter_patterns : list, optional + List of parameter patterns to obfuscate in the notebook. + Defaults to `utils.SENSITIVE_PARAMETER_PATTERNS` **kwargs Arbitrary keyword arguments to pass to the notebook engine @@ -102,6 +109,8 @@ def execute_notebook( kernel_name=kernel_name, language=language, engine_name=engine_name, + obfuscate_sensitive_parameters=obfuscate_sensitive_parameters, + sensitive_parameter_patterns=sensitive_parameter_patterns, ) nb = prepare_notebook_metadata(nb, input_path, output_path, report_mode) diff --git a/papermill/parameterize.py b/papermill/parameterize.py index b2422487..cd094f84 100644 --- a/papermill/parameterize.py +++ b/papermill/parameterize.py @@ -8,7 +8,7 @@ from .iorw import read_yaml_file from .log import logger from .translators import translate_parameters -from .utils import find_first_tagged_cell_index +from .utils import find_first_tagged_cell_index, obfuscate_parameters def add_builtin_parameters(parameters): @@ -64,6 +64,8 @@ def parameterize_notebook( kernel_name=None, language=None, engine_name=None, + obfuscate_sensitive_parameters=True, + sensitive_parameter_patterns=None, ): """Assigned parameters into the appropriate place in the input notebook @@ -77,6 +79,11 @@ def parameterize_notebook( Flag to set report mode comment : str, optional Comment added to the injected cell + obfuscate_sensitive_parameters : bool, optional + Obfuscate sensitive parameters in the notebook, Defaults to True + sensitive_parameter_patterns : list, optional + List of parameter patterns to obfuscate in the notebook. + Defaults to `utils.SENSITIVE_PARAMETER_PATTERNS` """ # Load from a file if 'parameters' is a string. if isinstance(parameters, str): @@ -95,6 +102,15 @@ def parameterize_notebook( newcell = nbformat.v4.new_code_cell(source=param_content) newcell.metadata['tags'] = ['injected-parameters'] + if obfuscate_sensitive_parameters: + obfuscated_param_content = translate_parameters( + kernel_name, + language, + obfuscate_parameters(parameters, sensitive_parameter_patterns), + comment, + ) + newcell.metadata['papermill-obfuscated-source'] = obfuscated_param_content + if report_mode: newcell.metadata['jupyter'] = newcell.get('jupyter', {}) newcell.metadata['jupyter']['source_hidden'] = True @@ -116,6 +132,8 @@ def parameterize_notebook( after = nb.cells nb.cells = before + [newcell] + after - nb.metadata.papermill['parameters'] = parameters + nb.metadata.papermill['parameters'] = parameters \ + if not obfuscate_sensitive_parameters \ + else obfuscate_parameters(parameters, sensitive_parameter_patterns) return nb diff --git a/papermill/tests/test_cli.py b/papermill/tests/test_cli.py index c893185b..ab28aa1a 100755 --- a/papermill/tests/test_cli.py +++ b/papermill/tests/test_cli.py @@ -89,6 +89,8 @@ class TestCLI(unittest.TestCase): execution_timeout=None, report_mode=False, cwd=None, + obfuscate_sensitive_parameters=True, + sensitive_parameter_patterns=None, stdout_file=None, stderr_file=None, ) @@ -536,3 +538,161 @@ def test_stdout_file(tmpdir): with open(str(stdout_file)) as fp: assert fp.read() == f"{secret}\n" + + +@require_papermill_installed +def test_obfuscated_output(): + metadata = {'kernelspec': {'name': 'python3', 'language': 'python', 'display_name': 'python3'}} + secret = str(uuid.uuid4()) + notebook = nbformat.v4.new_notebook( + metadata=metadata, + cells=[ + nbformat.v4.new_code_cell('print(safe_text, "=", token)'), + ], + ) + process = papermill_cli( + [ + '-', + '-', + '-p', + 'token', + secret, + '-p', + 'safe_text', + 'Test', + ], + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + text = nbformat.writes(notebook) + out, err = process.communicate(input=text.encode('utf-8')) + + # Test no message on std error + assert not err + + # Test that secrets in the output are obfuscated + output_notebook = nbformat.reads(out.decode('utf-8'), as_version=4) + + # secret in the notebook metadata should be obfuscated + assert output_notebook.metadata['papermill']['parameters'] == { + 'safe_text': 'Test', + 'token': '********', + } + + assert len(output_notebook.cells) == 2 + + # secret in the injected parameters cell should be obfuscated + assert output_notebook.cells[0].cell_type == 'code' + assert output_notebook.cells[0].source == '# Parameters\ntoken = "********"\nsafe_text = "Test"\n' + + # secret in the output cell will be printed as is + assert output_notebook.cells[1].cell_type == 'code' + assert len(output_notebook.cells[1].outputs) == 1 + assert output_notebook.cells[1].outputs[0].output_type == 'stream' + assert output_notebook.cells[1].outputs[0].text == f'Test = {secret}\n' + +@require_papermill_installed +def test_disable_output_obfuscation(): + metadata = {'kernelspec': {'name': 'python3', 'language': 'python', 'display_name': 'python3'}} + secret = str(uuid.uuid4()) + notebook = nbformat.v4.new_notebook( + metadata=metadata, + cells=[ + nbformat.v4.new_code_cell('print(safe_text, "=", token)'), + ], + ) + process = papermill_cli( + [ + '-', + '-', + '-p', + 'token', + secret, + '-p', + 'safe_text', + 'Test', + '--no-obfuscate-sensitive-parameters', + ], + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + text = nbformat.writes(notebook) + out, err = process.communicate(input=text.encode('utf-8')) + + # Test no message on std error + assert not err + + # Test that secrets in the output are not obfuscated + output_notebook = nbformat.reads(out.decode('utf-8'), as_version=4) + + # secret in the notebook metadata should not be obfuscated + assert output_notebook.metadata['papermill']['parameters'] == { + 'safe_text': 'Test', + 'token': secret, + } + + assert len(output_notebook.cells) == 2 + + # secret in the injected parameters cell should be obfuscated + assert output_notebook.cells[0].cell_type == 'code' + assert output_notebook.cells[0].source == f'# Parameters\ntoken = "{secret}"\nsafe_text = "Test"\n' + + # secret in the output cell will be printed as is + assert output_notebook.cells[1].cell_type == 'code' + assert len(output_notebook.cells[1].outputs) == 1 + assert output_notebook.cells[1].outputs[0].output_type == 'stream' + assert output_notebook.cells[1].outputs[0].text == f'Test = {secret}\n' + +@require_papermill_installed +def test_custom_output_obfuscation(): + metadata = {'kernelspec': {'name': 'python3', 'language': 'python', 'display_name': 'python3'}} + secret = str(uuid.uuid4()) + notebook = nbformat.v4.new_notebook( + metadata=metadata, + cells=[ + nbformat.v4.new_code_cell('print(safe_text, "=", token)'), + ], + ) + process = papermill_cli( + [ + '-', + '-', + '-p', + 'token', + secret, + '-p', + 'safe_text', + 'Test', + '--sensitive-parameter-patterns', + 'safe_text', + ], + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + text = nbformat.writes(notebook) + out, err = process.communicate(input=text.encode('utf-8')) + + # Test no message on std error + assert not err + + # Test that secrets in the output are obfuscated + output_notebook = nbformat.reads(out.decode('utf-8'), as_version=4) + + # secret in the notebook metadata should be obfuscated + # --sensitive-parameter-patterns should be set to obfuscate 'safe_text' and not 'token' + assert output_notebook.metadata['papermill']['parameters'] == { + 'safe_text': '********', + 'token': secret, + } + + assert len(output_notebook.cells) == 2 + + # secret in the injected parameters cell should be obfuscated + assert output_notebook.cells[0].cell_type == 'code' + assert output_notebook.cells[0].source == f'# Parameters\ntoken = "{secret}"\nsafe_text = "********"\n' + + # secret in the output cell will be printed as is + assert output_notebook.cells[1].cell_type == 'code' + assert len(output_notebook.cells[1].outputs) == 1 + assert output_notebook.cells[1].outputs[0].output_type == 'stream' + assert output_notebook.cells[1].outputs[0].text == f'Test = {secret}\n' diff --git a/papermill/tests/test_parameterize.py b/papermill/tests/test_parameterize.py index 431caa12..a9c390d8 100644 --- a/papermill/tests/test_parameterize.py +++ b/papermill/tests/test_parameterize.py @@ -81,6 +81,29 @@ def test_custom_comment(self): first_line = cell_one['source'].split('\n')[0] self.assertEqual(first_line, '# This is a custom comment') + def test_sensitive_parameters_obfuscation(self): + test_nb = load_notebook_node(get_notebook_path("simple_execute.ipynb")) + test_nb = parameterize_notebook(test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=True) + + cell_one = test_nb.cells[1] + self.assertIn('password', cell_one['source']) + self.assertIn('secret', cell_one['source']) + + self.assertIn('papermill-obfuscated-source', cell_one['metadata']) + obfuscated_output = cell_one['metadata']['papermill-obfuscated-source'] + self.assertIn('password', obfuscated_output) + self.assertNotIn('secret', obfuscated_output) + self.assertIn('********', obfuscated_output) + + def test_sensitive_parameters_obfuscation_disabled(self): + test_nb = load_notebook_node(get_notebook_path("simple_execute.ipynb")) + test_nb = parameterize_notebook(test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=False) + + cell_one = test_nb.cells[1] + self.assertIn('password', cell_one['source']) + self.assertIn('secret', cell_one['source']) + + self.assertNotIn('papermill-obfuscated-source', cell_one['metadata']) class TestBuiltinParameters(unittest.TestCase): def test_add_builtin_parameters_keeps_provided_parameters(self): diff --git a/papermill/tests/test_utils.py b/papermill/tests/test_utils.py index 4e24ce75..527eaf0c 100644 --- a/papermill/tests/test_utils.py +++ b/papermill/tests/test_utils.py @@ -13,6 +13,7 @@ merge_kwargs, remove_args, retry, + obfuscate_parameter, ) @@ -58,3 +59,32 @@ def test_chdir(): assert Path.cwd() == Path(temp_dir) assert Path.cwd() == old_cwd + + +def test_obfuscate_parameter(): + # *password* + assert obfuscate_parameter("password", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("sample_password", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("password_for_test", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("password", "") == "" + + # *token* + assert obfuscate_parameter("token", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("sample_token", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("token_for_test", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("token", "") == "" + + # *key* + assert obfuscate_parameter("key", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("sample_key", "string_to_be_obfuscated") == "********" + assert obfuscate_parameter("keyword", "string_not_to_be_obfuscated") == "string_not_to_be_obfuscated" + assert obfuscate_parameter("key", "") == "" + + +def test_obfuscate_parameter_custom_pattern(): + # *secret* + assert obfuscate_parameter("secret", "string_to_be_obfuscated", [".*secret"]) == "********" + assert obfuscate_parameter("sample_secret", "string_to_be_obfuscated", [".*secret"]) == "********" + assert obfuscate_parameter("secret_for_test", "string_to_be_obfuscated", [".*secret"]) == "********" + # If the custom pattern are set, the default pattern should not be applied + assert obfuscate_parameter("token", "string_to_be_obfuscated", [".*secret"]) == "string_to_be_obfuscated" diff --git a/papermill/utils.py b/papermill/utils.py index f7db55c1..2e62f035 100644 --- a/papermill/utils.py +++ b/papermill/utils.py @@ -1,5 +1,6 @@ import logging import os +import re import warnings from contextlib import contextmanager from functools import wraps @@ -9,6 +10,23 @@ logger = logging.getLogger('papermill.utils') +SENSITIVE_PARAMETER_PATTERNS = [ + r"(?i)pass(word|wd)", + r"(?i)pwd$", + # A short keyword that may match many keywords unintentionally + # is only targeted when placed at the end of the string. + r"(?i)pass$", + r"(?i)token", + r"(?i)secret", + r"(?i)authorization", + r"(?i)auth$", + r"(?i)key$", + r"(?i)access_key", + r"(?i)secret_key", + r"(?i)private_key", +] + + def any_tagged_cell(nb, tag): """Whether the notebook contains at least one cell tagged ``tag``? @@ -190,3 +208,61 @@ def chdir(path): yield finally: os.chdir(old_dir) + +def obfuscate_parameter( + name, + value, + name_patterns=None, + obfuscated_value="********", +): + """Obfuscate parameter if it is sensitive. + + Parameters + ---------- + name : str + The name of the parameter + value : str + The value of the parameter + name_patterns : list, optional + List of patterns to obfuscate in the notebook. Defaults to `DEFAULT_OBFUSCATE_PATTERNS` + obfuscated_value : str, optional + The value to replace the sensitive parameter with. Defaults to "********" + + Returns + ------- + str + The obfuscated value if the parameter is sensitive, otherwise the original value + """ + if name_patterns is None: + name_patterns = SENSITIVE_PARAMETER_PATTERNS + + if len(name_patterns) == 0: + raise ValueError("No name patterns provided to obfuscate") + if not value: + # Return empty string if value is empty to show that the parameter is empty + return value + if any(re.search(pattern, name) for pattern in name_patterns): + return obfuscated_value + return value + +def obfuscate_parameters(params, name_patterns=None, obfuscated_value="********"): + """Obfuscate parameters if they are sensitive. + + Parameters + ---------- + params : dict + The parameters to obfuscate + name_patterns : list, optional + List of patterns to obfuscate in the notebook. Defaults to `DEFAULT_OBFUSCATE_PATTERNS` + obfuscated_value : str, optional + The value to replace the sensitive parameter with. Defaults to "********" + + Returns + ------- + dict + The obfuscated parameters + """ + return { + name: obfuscate_parameter(name, value, name_patterns, obfuscated_value) + for name, value in params.items() + } From 2665986fc069bcfed6f7e791277783c8978391a3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Mar 2025 10:38:28 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- papermill/cli.py | 1 - papermill/execute.py | 2 +- papermill/parameterize.py | 6 ++++-- papermill/tests/test_cli.py | 2 ++ papermill/tests/test_parameterize.py | 9 +++++++-- papermill/tests/test_utils.py | 2 +- papermill/utils.py | 7 +++---- 7 files changed, 18 insertions(+), 11 deletions(-) diff --git a/papermill/cli.py b/papermill/cli.py index aca8d8f9..3ed345b5 100755 --- a/papermill/cli.py +++ b/papermill/cli.py @@ -15,7 +15,6 @@ from .execute import execute_notebook from .inspection import display_notebook_help from .iorw import NoDatesSafeLoader, read_yaml_file -from .utils import SENSITIVE_PARAMETER_PATTERNS from .version import version as papermill_version click.disable_unicode_literals_warning = True diff --git a/papermill/execute.py b/papermill/execute.py index 3674daf8..5dbb5a18 100644 --- a/papermill/execute.py +++ b/papermill/execute.py @@ -8,7 +8,7 @@ from .iorw import get_pretty_path, load_notebook_node, local_file_io_cwd, write_ipynb from .log import logger from .parameterize import add_builtin_parameters, parameterize_notebook, parameterize_path -from .utils import chdir, obfuscate_parameters +from .utils import chdir def execute_notebook( diff --git a/papermill/parameterize.py b/papermill/parameterize.py index cd094f84..2e03d25e 100644 --- a/papermill/parameterize.py +++ b/papermill/parameterize.py @@ -132,8 +132,10 @@ def parameterize_notebook( after = nb.cells nb.cells = before + [newcell] + after - nb.metadata.papermill['parameters'] = parameters \ - if not obfuscate_sensitive_parameters \ + nb.metadata.papermill['parameters'] = ( + parameters + if not obfuscate_sensitive_parameters else obfuscate_parameters(parameters, sensitive_parameter_patterns) + ) return nb diff --git a/papermill/tests/test_cli.py b/papermill/tests/test_cli.py index ab28aa1a..e9d46694 100755 --- a/papermill/tests/test_cli.py +++ b/papermill/tests/test_cli.py @@ -591,6 +591,7 @@ def test_obfuscated_output(): assert output_notebook.cells[1].outputs[0].output_type == 'stream' assert output_notebook.cells[1].outputs[0].text == f'Test = {secret}\n' + @require_papermill_installed def test_disable_output_obfuscation(): metadata = {'kernelspec': {'name': 'python3', 'language': 'python', 'display_name': 'python3'}} @@ -643,6 +644,7 @@ def test_disable_output_obfuscation(): assert output_notebook.cells[1].outputs[0].output_type == 'stream' assert output_notebook.cells[1].outputs[0].text == f'Test = {secret}\n' + @require_papermill_installed def test_custom_output_obfuscation(): metadata = {'kernelspec': {'name': 'python3', 'language': 'python', 'display_name': 'python3'}} diff --git a/papermill/tests/test_parameterize.py b/papermill/tests/test_parameterize.py index a9c390d8..9929940d 100644 --- a/papermill/tests/test_parameterize.py +++ b/papermill/tests/test_parameterize.py @@ -83,7 +83,9 @@ def test_custom_comment(self): def test_sensitive_parameters_obfuscation(self): test_nb = load_notebook_node(get_notebook_path("simple_execute.ipynb")) - test_nb = parameterize_notebook(test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=True) + test_nb = parameterize_notebook( + test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=True + ) cell_one = test_nb.cells[1] self.assertIn('password', cell_one['source']) @@ -97,7 +99,9 @@ def test_sensitive_parameters_obfuscation(self): def test_sensitive_parameters_obfuscation_disabled(self): test_nb = load_notebook_node(get_notebook_path("simple_execute.ipynb")) - test_nb = parameterize_notebook(test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=False) + test_nb = parameterize_notebook( + test_nb, {'msg': 'Hello', 'password': 'secret'}, obfuscate_sensitive_parameters=False + ) cell_one = test_nb.cells[1] self.assertIn('password', cell_one['source']) @@ -105,6 +109,7 @@ def test_sensitive_parameters_obfuscation_disabled(self): self.assertNotIn('papermill-obfuscated-source', cell_one['metadata']) + class TestBuiltinParameters(unittest.TestCase): def test_add_builtin_parameters_keeps_provided_parameters(self): with_builtin_parameters = add_builtin_parameters({"foo": "bar"}) diff --git a/papermill/tests/test_utils.py b/papermill/tests/test_utils.py index 527eaf0c..265edac8 100644 --- a/papermill/tests/test_utils.py +++ b/papermill/tests/test_utils.py @@ -11,9 +11,9 @@ any_tagged_cell, chdir, merge_kwargs, + obfuscate_parameter, remove_args, retry, - obfuscate_parameter, ) diff --git a/papermill/utils.py b/papermill/utils.py index 2e62f035..6f9d72c0 100644 --- a/papermill/utils.py +++ b/papermill/utils.py @@ -209,6 +209,7 @@ def chdir(path): finally: os.chdir(old_dir) + def obfuscate_parameter( name, value, @@ -245,6 +246,7 @@ def obfuscate_parameter( return obfuscated_value return value + def obfuscate_parameters(params, name_patterns=None, obfuscated_value="********"): """Obfuscate parameters if they are sensitive. @@ -262,7 +264,4 @@ def obfuscate_parameters(params, name_patterns=None, obfuscated_value="********" dict The obfuscated parameters """ - return { - name: obfuscate_parameter(name, value, name_patterns, obfuscated_value) - for name, value in params.items() - } + return {name: obfuscate_parameter(name, value, name_patterns, obfuscated_value) for name, value in params.items()}