Moved run_command into the class

Author: DeepHorizons

batchspawner/batchspawner.py

@@ -34,34 +34,6 @@
 from jupyterhub.spawner import set_user_setuid
 import jupyterhub
 
-@gen.coroutine
-def run_command(cmd, input=None, env=None):
-    proc = Subprocess(cmd, shell=True, env=env, stdin=Subprocess.STREAM, stdout=Subprocess.STREAM,stderr=Subprocess.STREAM)
-    inbytes = None
-    if input:
-        inbytes = input.encode()
-        try:
-            yield proc.stdin.write(inbytes)
-        except StreamClosedError as exp:
-            # Apparently harmless
-            pass
-    proc.stdin.close()
-    out = yield proc.stdout.read_until_close()
-    eout = yield proc.stderr.read_until_close()
-    proc.stdout.close()
-    proc.stderr.close()
-    eout = eout.decode().strip()
-    try:
-        err = yield proc.wait_for_exit()
-    except CalledProcessError:
-        #self.log.error("Subprocess returned exitcode %s" % proc.returncode)
-        #self.log.error(eout)
-        raise RuntimeError(eout)
-    if err != 0:
-        return err # exit error?
-    else:
-        out = out.decode().strip()
-        return out
 
 class BatchSpawnerBase(Spawner):
     """Base class for spawners using resource manager batch job submission mechanisms
@@ -174,6 +146,35 @@ def parse_job_id(self, output):
     def cmd_formatted_for_batch(self):
         return ' '.join(self.cmd + self.get_args())
 
+    @gen.coroutine
+    def run_command(self, cmd, input=None, env=None):
+        proc = Subprocess(cmd, shell=True, env=env, stdin=Subprocess.STREAM, stdout=Subprocess.STREAM,stderr=Subprocess.STREAM)
+        inbytes = None
+        if input:
+            inbytes = input.encode()
+            try:
+                yield proc.stdin.write(inbytes)
+            except StreamClosedError as exp:
+                # Apparently harmless
+                pass
+        proc.stdin.close()
+        out = yield proc.stdout.read_until_close()
+        eout = yield proc.stderr.read_until_close()
+        proc.stdout.close()
+        proc.stderr.close()
+        eout = eout.decode().strip()
+        try:
+            err = yield proc.wait_for_exit()
+        except CalledProcessError:
+            #self.log.error("Subprocess returned exitcode %s" % proc.returncode)
+            #self.log.error(eout)
+            raise RuntimeError(eout)
+        if err != 0:
+            return err # exit error?
+        else:
+            out = out.decode().strip()
+            return out
+
     @gen.coroutine
     def submit_batch_script(self):
         subvars = self.get_req_subvars()
@@ -184,7 +185,7 @@ def submit_batch_script(self):
         script = self.batch_script.format(**subvars)
         self.log.info('Spawner submitting job using ' + cmd)
         self.log.info('Spawner submitted script:\n' + script)
-        out = yield run_command(cmd, input=script, env=self.get_env())
+        out = yield self.run_command(cmd, input=script, env=self.get_env())
         try:
             self.log.info('Job submitted. cmd: ' + cmd + ' output: ' + out)
             self.job_id = self.parse_job_id(out)
@@ -210,7 +211,7 @@ def read_job_state(self):
         cmd = self.batch_query_cmd.format(**subvars)
         self.log.debug('Spawner querying job: ' + cmd)
         try:
-            out = yield run_command(cmd)
+            out = yield self.run_command(cmd)
             self.job_status = out
         except Exception as e:
             self.log.error('Error querying job ' + self.job_id)
@@ -228,7 +229,7 @@ def cancel_batch_job(self):
         subvars['job_id'] = self.job_id
         cmd = self.batch_cancel_cmd.format(**subvars)
         self.log.info('Cancelling job ' + self.job_id + ': ' + cmd)
-        yield run_command(cmd)
+        yield self.run_command(cmd)
 
     def load_state(self, state):
         """load job_id from state"""