Python wrapper for Jupyter Notebooks Public API

dataikuapi/dss/notebook.py

@@ -16,44 +16,46 @@ def unload(self, session_id=None):
         """
         Stop this Jupyter notebook and release its resources
         """
-        state = self.get_state()
-        if state is None:
+        sessions = self.get_sessions()
+        if sessions is None:
             raise Exception("Notebook isn't running")
-        if state.get('activeSessions', None) is None:
-            raise Exception("Notebook isn't running")
-        if len(state['activeSessions']) == 0:
+        if len(sessions) == 0:
             raise Exception("Notebook isn't running")
         if session_id is None:
-            if len(state['activeSessions']) > 1:
+            if len(sessions) > 1:
                 raise Exception("Several sessions of the notebook are running, choose one")
             else:
-                session_id = state['activeSessions'][0].get('sessionId', None)
+                session_id = sessions[0].get('sessionId', None)
         return self.client._perform_json("DELETE",
                                          "/projects/%s/jupyter-notebooks/%s/sessions/%s" % (self.project_key, self.notebook_name, session_id))
 
-    def get_state(self):
+    def get_state(self, refresh=False):
         """
         Get the status of this Jupyter notebook
+
+        :param bool refresh: if True, get the status of the notebook from the backend
         """
         notebook_states = self.client._perform_json("GET",
-                                                 "/projects/%s/jupyter-notebooks/" % self.project_key,
-                                                 params={"active": False})
-        for notebook in notebook_states:
-            if notebook.get("name") == self.notebook_name:
-                    self.state = notebook
+                                                    "/projects/%s/jupyter-notebooks/" % self.project_key,
+                                                    params={"active": False})
+        if self.state is None or refresh:
+            for state in notebook_states:
+                if state.get("name") == self.notebook_name:
+                    self.state = state
                     return self.state
         return self.state
 
     def get_sessions(self):
         """
-        Get the list of the running sessions of this Jupyter notebook
+        Get the list of running sessions of this Jupyter notebook
         """
-        state = self.get_state()
-        if state is None:
-            raise Exception("Notebook isn't running")
-        if state.get('activeSessions', None) is None:
-            raise Exception("Notebook isn't running")
-        return state['activeSessions']
+
+        if self.state is None:
+            self.state = {}
+        sessions = self.client._perform_json("GET",
+                                             "/projects/%s/jupyter-notebooks/%s/sessions" % (self.project_key, self.notebook_name))
+        self.state["activeSessions"] = sessions
+        return sessions
 
     def get_content(self):
         """

dataikuapi/dss/project.py

@@ -862,16 +862,19 @@ def get_jupyter_notebook(self, notebook_name):
 
     def create_jupyter_notebook(self, notebook_name, notebook_content):
         """
-        Create a new jupyter notebook.
+        Create a new jupyter notebook and get a handle to interact with it
 
         :param str notebook_name: the name of the notebook to create
         :param dict notebook_content: the data of the notebook to create, as a dict.
             The data will be converted to a JSON string internally.
             Use ``get_state()`` on a similar existing ``DSSNotebook`` object in order to get a sample definition object
+        :returns: A handle to interact with the newly created jupyter notebook
+        :rtype: :class:`~dataikuapi.dss.notebook.DSSNotebook` jupyter notebook handle
         """
-        return self.client._perform_json("POST",
-                                         "/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name),
-                                         body=notebook_content)
+        self.client._perform_json("POST",
+                                  "/projects/%s/jupyter-notebooks/%s" % (self.project_key, notebook_name),
+                                  body=notebook_content)
+        return DSSNotebook(self.client, self.project_key, notebook_name, content=notebook_content)
 
     ########################################################
     # Continuous activities