Copy parameter cell metadata to injected cell

Author: paranoa233

papermill/parameterize.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from datetime import datetime, timezone
 from uuid import uuid4
 
@@ -92,15 +93,21 @@ def parameterize_notebook(
     # Upgrade the Notebook to the latest v4 before writing into it
     nb = nbformat.v4.upgrade(nb)
 
+    param_cell_index = find_first_tagged_cell_index(nb, 'parameters')
+    injected_cell_index = find_first_tagged_cell_index(nb, 'injected-parameters')
+
     newcell = nbformat.v4.new_code_cell(source=param_content)
-    newcell.metadata['tags'] = ['injected-parameters']
+    if param_cell_index >= 0:
+        newcell.metadata = deepcopy(nb.cells[param_cell_index].metadata)
+
+    tags = [tag for tag in newcell.metadata.get('tags', []) if tag != 'parameters']
+    if 'injected-parameters' not in tags:
+        tags.append('injected-parameters')
+    newcell.metadata['tags'] = tags
 
     if report_mode:
-        newcell.metadata['jupyter'] = newcell.get('jupyter', {})
-        newcell.metadata['jupyter']['source_hidden'] = True
+        newcell.metadata.setdefault('jupyter', {})['source_hidden'] = True
 
-    param_cell_index = find_first_tagged_cell_index(nb, 'parameters')
-    injected_cell_index = find_first_tagged_cell_index(nb, 'injected-parameters')
     if injected_cell_index >= 0:
         # Replace the injected cell with a new version
         before = nb.cells[:injected_cell_index]

papermill/tests/test_parameterize.py

@@ -11,20 +11,23 @@ class TestNotebookParametrizing(unittest.TestCase):
     def count_nb_injected_parameter_cells(self, nb):
         return len([c for c in nb.cells if 'injected-parameters' in c.get('metadata', {}).get('tags', [])])
 
-    def test_no_tag_copying(self):
-        # Test that injected cell does not copy other tags
+    def test_parameter_cell_metadata_is_copied(self):
         test_nb = load_notebook_node(get_notebook_path("simple_execute.ipynb"))
         test_nb.cells[0]['metadata']['tags'].append('some tag')
+        test_nb.cells[0]['metadata']['slideshow'] = {'slide_type': 'skip'}
 
         test_nb = parameterize_notebook(test_nb, {'msg': 'Hello'})
 
         cell_zero = test_nb.cells[0]
         self.assertTrue('some tag' in cell_zero.get('metadata').get('tags'))
         self.assertTrue('parameters' in cell_zero.get('metadata').get('tags'))
+        self.assertEqual({'slide_type': 'skip'}, cell_zero.get('metadata').get('slideshow'))
 
         cell_one = test_nb.cells[1]
-        self.assertTrue('some tag' not in cell_one.get('metadata').get('tags'))
+        self.assertTrue('some tag' in cell_one.get('metadata').get('tags'))
+        self.assertTrue('parameters' not in cell_one.get('metadata').get('tags'))
         self.assertTrue('injected-parameters' in cell_one.get('metadata').get('tags'))
+        self.assertEqual({'slide_type': 'skip'}, cell_one.get('metadata').get('slideshow'))
 
         self.assertEqual(self.count_nb_injected_parameter_cells(test_nb), 1)