Added cell description to tqdm

papermill/engines.py

@@ -213,6 +213,11 @@ def cell_start(self, cell, cell_index=None, **kwargs):
         cell.metadata.papermill["status"] = self.RUNNING
         cell.metadata.papermill['exception'] = False
 
+        # injects optional description of the current cell directly in the tqdm
+        cell_description = self.get_cell_description(cell)
+        if cell_description is not None and hasattr(self, 'pbar') and self.pbar:
+            self.pbar.set_description(f"Executing {cell_description}")
+
         self.save()
 
     @catch_nb_assignment
@@ -284,6 +289,17 @@ def notebook_complete(self, **kwargs):
         # Force a final sync
         self.save()
 
+    def get_cell_description(self, cell, escape_str="papermill_description="):
+        """Fetches cell description if present"""
+        if cell is None:
+            return None
+
+        cell_code = cell["source"]
+        if cell_code is None or escape_str not in cell_code:
+            return None
+
+        return cell_code.split(escape_str)[1].split()[0]
+
     def complete_pbar(self):
         """Refresh progress bar"""
         if hasattr(self, 'pbar') and self.pbar:

papermill/tests/notebooks/simple_execute.ipynb

@@ -19,6 +19,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "#papermill_description=DESC\n",
     "print(msg)"
    ]
   },
@@ -46,9 +47,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.8.5"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }

papermill/tests/test_engines.py

@@ -81,10 +81,15 @@ def test_save_new_nb(self):
         nb_man = NotebookExecutionManager(self.nb)
         nb_man.save(nb=self.foo_nb)
         self.assertEqual(nb_man.nb.metadata['foo'], 'bar')
-
+
+    def test_get_cell_description(self):
+        nb_man = NotebookExecutionManager(self.nb)
+        self.assertIsNone(nb_man.get_cell_description(nb_man.nb.cells[0]))
+        self.assertEqual(nb_man.get_cell_description(nb_man.nb.cells[1]), 'DESC')
+
     def test_notebook_start(self):
         nb_man = NotebookExecutionManager(self.nb)
-        nb_man.save = Mock()
+        nb_man.nb_path.metadata['foo'] = 'bar'
         nb_man.notebook_start()
 
         self.assertEqual(nb_man.nb.metadata.papermill['start_time'], nb_man.start_time.isoformat())