diff --git a/appveyor.yml b/appveyor.yml index b9ba789..c2969f2 100644 --- a/appveyor.yml +++ b/appveyor.yml @@ -13,8 +13,9 @@ install: - conda config --set always_yes yes --set changeps1 no - conda update -q conda - conda info -a - - conda create -q -n test-environment python=%PYTHON_VERSION% ipython pytest + - conda create -q -n test-environment python=%PYTHON_VERSION% ipython pytest nbformat nbconvert ipykernel - activate test-environment + - python -m ipykernel install --user --name python3 --display-name "Python 3" test_script: - set PYTHONPATH=%PYTHONPATH%;%CD% diff --git a/watermark/magic.py b/watermark/magic.py index 01dd9b5..40e6271 100644 --- a/watermark/magic.py +++ b/watermark/magic.py @@ -79,6 +79,8 @@ class WaterMark(Magics): help='include information about how Python was installed') @argument('--check_latest', action='store_true', help='check if the latest packages are installed') + @argument("-md", "--metadata", action="store_true", + help="save watermark info to notebook metadata") @line_magic def watermark(self, line): """ diff --git a/watermark/tests/test_html_export.py b/watermark/tests/test_html_export.py new file mode 100644 index 0000000..48451be --- /dev/null +++ b/watermark/tests/test_html_export.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +import nbformat +from nbconvert.preprocessors import ExecutePreprocessor +import unittest + +def test_html_output_presence(): + + nb = nbformat.v4.new_notebook() + + code = "%load_ext watermark\n%watermark" + nb.cells.append(nbformat.v4.new_code_cell(code)) + + ep = ExecutePreprocessor(timeout=600, kernel_name='python3') + + try: + ep.preprocess(nb, {'metadata': {'path': '.'}}) + + cell_output = nb.cells[0].outputs[0].text + + assert 'Python implementation' in cell_output + assert 'IPython version' in cell_output + print("\n HTML Output Test Passed!") + + except Exception as e: + print(f"\n Test failed due to: {e}") + raise e + +if __name__ == "__main__": + test_html_output_presence() \ No newline at end of file diff --git a/watermark/watermark.py b/watermark/watermark.py index 1950bce..b46370e 100644 --- a/watermark/watermark.py +++ b/watermark/watermark.py @@ -61,6 +61,7 @@ def watermark( gpu=False, jupyter_env=False, python_installation=False, + metadata=False, check_latest=False, watermark_self=None, globals_=None @@ -231,6 +232,43 @@ def watermark( output.append({"Jupyter enviroment": _get_jupyter_env()}) if args['watermark']: output.append({"Watermark": __version__}) + if args.get('metadata'): + try: + flat_output = {} + for d in output: + if d: + flat_output.update(d) + + from IPython import get_ipython + import nbformat + import os + + shell = get_ipython() + if 'filename' in shell.user_ns: + nb_path = shell.user_ns['filename'] + else: + import ipykernel + from notebook.notebookapp import list_running_servers + pass + + from IPython.display import display, Javascript + import json + + js_payload = json.dumps(flat_output) + js_code = f""" + (function() {{ + if (window.IPython && IPython.notebook) {{ + IPython.notebook.metadata.watermark = {js_payload}; + console.log("Watermark metadata saved to notebook."); + }} else {{ + console.warn("Watermark: Notebook metadata storage is only supported in Classic Jupyter Notebook."); + }} + }})(); + """ + display(Javascript(js_code)) + + except Exception: + pass return _generate_formatted_text(output) @@ -472,3 +510,4 @@ def _get_latest_version(package_name): return data['info']['version'] except Exception: return None + \ No newline at end of file