From 1ad5ed59bb7b73586d8bda892abc9850b8194e10 Mon Sep 17 00:00:00 2001
From: Richard Darst <rkd@zgib.net>
Date: Wed, 2 May 2018 11:13:56 +0300
Subject: [PATCH] Add progress indicators when spawning

- Uses JupyterHub 0.9 feature, no effect for <0.9.
- Closes: #81
---
 .travis.yml                  |  1 +
 batchspawner/batchspawner.py | 31 +++++++++++++++++++++++++++++++
 requirements.txt             |  1 +
 3 files changed, 33 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index 343b4ec4..46e66f44 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -35,6 +35,7 @@ install:
     - if [ $JHUB_VER != "master" -a $JHUB_VER != "0.9.0b2" ]; then pip install "tornado<5.0"; fi
     - pip install --pre -f travis-wheels/wheelhouse -r jupyterhub/dev-requirements.txt
     - pip install --pre -e jupyterhub
+    - pip install --pre -f travis-wheels/wheelhouse -r requirements.txt
 
 script:
     - travis_retry py.test --lf --cov batchspawner batchspawner/tests -v
diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py
index 5dde6bf1..bb50446d 100644
--- a/batchspawner/batchspawner.py
+++ b/batchspawner/batchspawner.py
@@ -17,9 +17,16 @@
 """
 import pwd
 import os
+import sys
 
 import xml.etree.ElementTree as ET
 
+# For the progress method of JupyterHub 0.9, requiring Python 3.5+.
+# If not supported, don't support it which is OK.
+if sys.version_info >= (3, 5):
+    import asyncio
+    from async_generator import async_generator, yield_, yield_from_
+
 from tornado import gen
 from tornado.process import Subprocess
 from subprocess import CalledProcessError
@@ -353,6 +360,28 @@ def stop(self, now=False):
                              self.job_id, self.current_ip, self.port)
                 )
 
+    # The progress method is only supported on JupyterHub 0.9, which
+    # only supports Python 3.5+.  It requires asyncio coroutines.
+    if sys.version_info >= (3, 5):
+        @async_generator
+        @asyncio.coroutine
+        def progress(self):
+            while True:
+                if self.state_ispending():
+                    yield_from_({
+                        "message": "Pending in queue...",
+                    })
+                elif self.state_isrunning():
+                    yield_from_({
+                        "message": "Cluster job running... waiting to connect",
+                    })
+                    return
+                else:
+                    yield_from_({
+                        "message": "Unknown status...",
+                    })
+                yield_from_(gen.sleep(.1))
+
 import re
 
 class BatchSpawnerRegexStates(BatchSpawnerBase):
@@ -486,6 +515,8 @@ class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
 #SBATCH --get-user-env=L
 #SBATCH {options}
 
+sleep 30
+
 which jupyterhub-singleuser
 {cmd}
 """).tag(config=True)
diff --git a/requirements.txt b/requirements.txt
index 19ec4dee..82a85e0a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
+async_generator>=1.8; python_version >= '3.5'
 jupyterhub>=0.5