From 2abb7c305f700518b9053ad548349573e2f583a8 Mon Sep 17 00:00:00 2001
From: Daniel Cizin <daniel.cizin@gmail.com>
Date: Wed, 27 Oct 2021 14:51:29 -0400
Subject: [PATCH] [v1.0.51] added jupyter magic support (#51)

* added jupyter magic support %%notify

* Add the %%notify Jupyter command

* Add Jupyter magic command docs

* updated readme

Co-authored-by: Cizin <cizind@lski3119.mskcc.root.mskcc.org>
Co-authored-by: Snir Shechter <snir@aporia.com>
---
 README.md                         | 23 +++++++++++++++++++++++
 sdk/pyproject.toml                |  2 +-
 sdk/src/mlnotify/__init__.py      |  9 +++++++++
 sdk/src/mlnotify/jupyter_magic.py | 19 +++++++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 sdk/src/mlnotify/jupyter_magic.py

diff --git a/README.md b/README.md
index 77c0989..3e785cb 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,7 @@ my_model.fit(...)
 ```
 
 The import will automagically hook into your fit/train method.
+
 - Once you start training your model a tracking url & QR code will be printed to the console.
   ![Printed tracking URL & QR code](docs/assets/printed-tracking-url-and-qr-code.png)
 - Enter the tracking url
@@ -73,6 +74,22 @@ Supported ML frameworks:
 - [Tensorflow](https://www.tensorflow.org/)
 - [Catboost](https://catboost.ai/)
 
+### Jupyter Notebook
+
+`import mlnotify`
+
+And in any Jupyter cell:
+
+```
+%%notify
+...
+```
+
+Works with line magic, too
+```
+%notify your_code()
+```
+
 ### Manual
 
 The library also exports a manual API to be used if you want to do it manually.
@@ -111,12 +128,15 @@ model.fit(...)
 ## API
 
 The library exports four items:
+
 ```python
 from mlnotify import start, end, plugins_manager, BasePlugin
 ```
+
 ### `start() -> None`
 
 Starts tracking.
+
 ### `end() -> None`
 
 Ends tracking.
@@ -135,10 +155,13 @@ Methods:
   Removes all registered plugins.
 
 ## Security
+
 No sensitive data is sent to the MLNotify server - only training start & end time.
 
 ## Contribution, self-deployment & local development
+
 Contributions and self-deployments are more than welcome.
+
 ### Website & API
 
 This project relies heavily on SaaS products, and must receive proper config for Netlify, Firebase and SendGrid for it to work. You can run this project locally using the Netlify CLI:
diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml
index 4f2508e..be4c7f6 100644
--- a/sdk/pyproject.toml
+++ b/sdk/pyproject.toml
@@ -4,7 +4,7 @@ description = "No need to keep checking your training. Add just 1 import line an
 homepage = 'https://mlnotify.aporia.com'
 name = "mlnotify"
 repository = "https://github.com/aporia-ai/mlnotify"
-version = "v1.0.50"
+version = "v1.0.51"
 
 [tool.poetry.dependencies]
 python = "^3.6"
diff --git a/sdk/src/mlnotify/__init__.py b/sdk/src/mlnotify/__init__.py
index cb3411f..79bd72c 100644
--- a/sdk/src/mlnotify/__init__.py
+++ b/sdk/src/mlnotify/__init__.py
@@ -4,4 +4,13 @@
 start = plugin_manager.run_before
 end = plugin_manager.run_after
 
+
+try:
+    from mlnotify.jupyter_magic import register_jupyter_magic
+    register_jupyter_magic()
+except:
+    # Not in jupyter notebook
+    pass
+
+
 __all__ = [start, end, plugin_manager, BasePlugin]
diff --git a/sdk/src/mlnotify/jupyter_magic.py b/sdk/src/mlnotify/jupyter_magic.py
new file mode 100644
index 0000000..55d781a
--- /dev/null
+++ b/sdk/src/mlnotify/jupyter_magic.py
@@ -0,0 +1,19 @@
+from IPython.core.magic import Magics, magics_class, line_cell_magic
+from mlnotify.mlnotify import plugin_manager
+
+# Jupyter line and cell magic
+@magics_class
+class MLNotifyMagic(Magics):
+    @line_cell_magic
+    def notify(self, line, cell=None):
+        plugin_manager.run_before()
+        self.shell.run_cell(line)
+        if cell is not None:
+            self.shell.run_cell(cell)
+        plugin_manager.run_after()
+
+
+def register_jupyter_magic():
+    ipython = get_ipython()
+    ipython.register_magics(MLNotifyMagic)
+