/scripts/exprlib/subproc.py

https://github.com/ANLAB-KAIST/NBA · Python · 63 lines · 46 code · 2 blank · 15 comment · 0 complexity · 048f058e5143f3dbc37c498dfb200262 MD5 · raw file

  1. #! /usr/bin/env python3
  2. import sys, os
  3. import asyncio
  4. import subprocess
  5. def execute(cmdargs, shell=False, iterable=False, async=False, read=False):
  6. '''
  7. Executes a shell command or external program using argument list.
  8. The output can be read after blocking or returned as iterable for further
  9. processing.
  10. This implementation is based on Snakemake's shell module.
  11. '''
  12. def _iter_stdout(proc, cmd):
  13. for line in proc.stdout:
  14. yield line[:-1] # strip newline at the end
  15. retcode = proc.wait()
  16. if retcode != 0:
  17. raise subprocess.CalledProcessError(retcode, cmd)
  18. popen_args = {}
  19. if shell and 'SHELL' in os.environ:
  20. popen_args['executable'] = os.environ['SHELL']
  21. close_fds = (sys.platform != 'win32')
  22. stdout = subprocess.PIPE if (iterable or async or read) else sys.stdout
  23. proc = subprocess.Popen(cmdargs, shell=shell, stdout=stdout,
  24. close_fds=close_fds, **popen_args)
  25. ret = None
  26. if iterable:
  27. return _iter_stdout(proc, cmdargs)
  28. if read:
  29. ret = proc.stdout.read().decode('utf8')
  30. elif async:
  31. return proc
  32. retcode = proc.wait()
  33. if retcode != 0:
  34. raise subprocess.CalledProcessError(retcode, cmdargs)
  35. return ret
  36. @asyncio.coroutine
  37. def execute_async_simple(cmdargs, timeout=None):
  38. proc = yield from asyncio.create_subprocess_exec(*cmdargs, stdout=sys.stdout)
  39. if timeout:
  40. try:
  41. retcode = yield from asyncio.wait_for(proc.wait(), timeout + 1)
  42. except asyncio.TimeoutError:
  43. print('Terminating the main process...', file=sys.stderr)
  44. proc.send_signal(signal.SIGINT)
  45. retcode = 0
  46. else:
  47. retcode = yield from proc.wait()
  48. return retcode
  49. _execute_memo = {}
  50. def execute_memoized(cmd):
  51. if cmd in _execute_memo:
  52. return _execute_memo[cmd]
  53. else:
  54. ret = execute(cmd, shell=True, read=True)
  55. _execute_memo[cmd] = ret
  56. return ret