From 458c817d83449f4c806c73cb93de60d2ee953c13 Mon Sep 17 00:00:00 2001
From: qimingj <qimingj@users.noreply.github.com>
Date: Tue, 6 Nov 2018 10:55:56 -0800
Subject: [PATCH] Add %%docker magic to jupyter kernel. It helps submitting a
 docker build job more easily with one cell. (#72)

---
 sdk/python/kfp/notebook/__init__.py | 16 +++++++++++
 sdk/python/kfp/notebook/_magic.py   | 42 +++++++++++++++++++++++++++++
 sdk/python/setup.py                 |  1 +
 3 files changed, 59 insertions(+)
 create mode 100644 sdk/python/kfp/notebook/__init__.py
 create mode 100644 sdk/python/kfp/notebook/_magic.py

diff --git a/sdk/python/kfp/notebook/__init__.py b/sdk/python/kfp/notebook/__init__.py
new file mode 100644
index 00000000000..abe647f1822
--- /dev/null
+++ b/sdk/python/kfp/notebook/__init__.py
@@ -0,0 +1,16 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from . import _magic
diff --git a/sdk/python/kfp/notebook/_magic.py b/sdk/python/kfp/notebook/_magic.py
new file mode 100644
index 00000000000..a1bb1b5ce92
--- /dev/null
+++ b/sdk/python/kfp/notebook/_magic.py
@@ -0,0 +1,42 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+try:
+  import IPython
+  import IPython.core.magic
+except ImportError:
+  raise Exception('This module can only be loaded in Jupyter.')
+
+
+from kfp.compiler import build_docker_image
+import tempfile
+
+
+@IPython.core.magic.register_cell_magic
+def docker(line, cell):
+  """cell magic for %%docker"""
+
+  if len(line.split()) != 2:
+    raise ValueError("usage: %%docker [gcr.io/project/image:tag] [gs://staging-bucket]")
+  if not cell.strip():
+    raise ValueError("Please fill in a dockerfile content in the cell.")
+
+  target, staging = line.split()
+  
+  with tempfile.NamedTemporaryFile(mode='wt') as f:
+    f.write(cell)
+    build_docker_image(staging, target, f.name)
+
+  
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
index 6f2ba6e00b5..a17fb7cc011 100644
--- a/sdk/python/setup.py
+++ b/sdk/python/setup.py
@@ -32,6 +32,7 @@
       'kfp.compiler',
       'kfp.components',
       'kfp.dsl',
+      'kfp.notebook',
       'kfp_experiment',
       'kfp_experiment.api',
       'kfp_experiment.models',