adds first implementation of NotebookIntrospector.get_params()

src/sklearn_evaluation/nb/NotebookIntrospector.py

@@ -73,6 +73,26 @@ def _parse_output(output, literal_eval, to_df, text_only):
                                                                to_df=to_df)
 
 
+def find_cell_with_tag(cells):
+    for cell in cells:
+        if ('metadata' in cell and 'tags' in cell['metadata']
+                and 'parameters' in cell['metadata']['tags']):
+            return cell
+
+
+def parse_parameters_cell(cells):
+    # this is a very simple implementation, for a more robust solution
+    # re-implement with ast or parso
+    cell = find_cell_with_tag(cells)
+
+    if not cell:
+        return dict()
+
+    tuples = [line.split('=') for line in cell['source'].splitlines()]
+
+    return {t[0].strip(): ast.literal_eval(t[1].strip()) for t in tuples}
+
+
 class NotebookIntrospector(Mapping):
     """Retrieve output from a notebook file with tagged cells.
 
@@ -145,3 +165,6 @@ def to_json_serializable(self):
                              text_only=True)
             for k, v in self.tag2output_raw.items()
         }
+
+    def get_params(self):
+        return parse_parameters_cell(self.nb.cells)

tests/nb/conftest.py

@@ -101,3 +101,16 @@ def nb_invalid_output():
 np.array([1, 2, 3])
 """
     return save_and_execute_notebook(content, 'nb_invalid_output.ipynb')
+
+
+@pytest.fixture
+def nb_parameters():
+    content = """
+import numpy as np
+
+# + tags=["parameters"]
+x = 1
+y = [1, 2]
+z = {'a': 1, 'b': 2}
+"""
+    return save_and_execute_notebook(content, 'nb_parameters.ipynb')

tests/nb/test_NotebookIntrospector.py

@@ -1,3 +1,5 @@
+import ast
+
 from sklearn_evaluation import NotebookIntrospector
 from IPython.display import HTML, Image
 
@@ -50,3 +52,19 @@ def test_json_serializable(tmp_directory, nb_plot):
     # "<Figure size YxZ with 1 Axes>")
     assert 'Figure size' in d['a']
     assert d['b'] == 42
+
+
+def test_get_params(tmp_directory, nb_parameters):
+    d = NotebookIntrospector('nb_parameters.ipynb')
+    assert d.get_params() == {
+        'x': 1,
+        'y': [1, 2],
+        'z': {
+            'a': 1,
+            'b': 2
+        },
+        'z': {
+            'a': 1,
+            'b': 2
+        }
+    }