You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

105 lines
3.6 KiB

  1. import os
  2. import re
  3. import subprocess
  4. import sys
  5. from functools import partial
  6. from threading import Thread
  7. from gooey.gui import events
  8. from gooey.gui.pubsub import pub
  9. from gooey.gui.util.casting import safe_float
  10. from gooey.gui.util.taskkill import taskkill
  11. from gooey.util.functional import unit, bind
  12. class ProcessController(object):
  13. def __init__(self, progress_regex, progress_expr, encoding, shell=True):
  14. self._process = None
  15. self.progress_regex = progress_regex
  16. self.progress_expr = progress_expr
  17. self.encoding = encoding
  18. self.wasForcefullyStopped = False
  19. self.shell_execution = shell
  20. def was_success(self):
  21. self._process.communicate()
  22. return self._process.returncode == 0
  23. def poll(self):
  24. if not self._process:
  25. raise Exception('Not started!')
  26. self._process.poll()
  27. def stop(self):
  28. if self.running():
  29. self.wasForcefullyStopped = True
  30. taskkill(self._process.pid)
  31. def running(self):
  32. return self._process and self.poll() is None
  33. def run(self, command):
  34. self.wasForcefullyStopped = False
  35. env = os.environ.copy()
  36. env["GOOEY"] = "1"
  37. env["PYTHONIOENCODING"] = self.encoding
  38. try:
  39. self._process = subprocess.Popen(
  40. command.encode(sys.getfilesystemencoding()),
  41. bufsize=1, stdout=subprocess.PIPE, stdin=subprocess.PIPE,
  42. stderr=subprocess.STDOUT, shell=self.shell_execution, env=env)
  43. except:
  44. self._process = subprocess.Popen(
  45. command,
  46. bufsize=1, stdout=subprocess.PIPE, stdin=subprocess.PIPE,
  47. stderr = subprocess.STDOUT, shell = self.shell_execution, env=env)
  48. t = Thread(target=self._forward_stdout, args=(self._process,))
  49. t.start()
  50. def _forward_stdout(self, process):
  51. '''
  52. Reads the stdout of `process` and forwards lines and progress
  53. to any interested subscribers
  54. '''
  55. while True:
  56. line = process.stdout.readline()
  57. if not line:
  58. break
  59. pub.send_message(events.CONSOLE_UPDATE, msg=line.decode(self.encoding))
  60. pub.send_message(events.PROGRESS_UPDATE,
  61. progress=self._extract_progress(line))
  62. pub.send_message(events.EXECUTION_COMPLETE)
  63. def _extract_progress(self, text):
  64. '''
  65. Finds progress information in the text using the
  66. user-supplied regex and calculation instructions
  67. '''
  68. # monad-ish dispatch to avoid the if/else soup
  69. find = partial(re.search, string=text.strip().decode(self.encoding))
  70. regex = unit(self.progress_regex)
  71. match = bind(regex, find)
  72. result = bind(match, self._calculate_progress)
  73. return result
  74. def _calculate_progress(self, match):
  75. '''
  76. Calculates the final progress value found by the regex
  77. '''
  78. if not self.progress_expr:
  79. return safe_float(match.group(1))
  80. else:
  81. return self._eval_progress(match)
  82. def _eval_progress(self, match):
  83. '''
  84. Runs the user-supplied progress calculation rule
  85. '''
  86. _locals = {k: safe_float(v) for k, v in match.groupdict().items()}
  87. if "x" not in _locals:
  88. _locals["x"] = [safe_float(x) for x in match.groups()]
  89. try:
  90. return int(eval(self.progress_expr, {}, _locals))
  91. except:
  92. return None