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