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.
 
 
 

100 lines
3.1 KiB

  1. import Queue
  2. import threading
  3. import sys
  4. import contextlib
  5. # This file provides a context manager that converts a function
  6. # that takes a callback into a generator that returns an iterable.
  7. # This is done by running the function in a new thread.
  8. # Based partially on http://stackoverflow.com/questions/9968592/
  9. class IteratorizerThread(threading.Thread):
  10. def __init__(self, queue, function, curl_hack):
  11. """
  12. function: function to execute, which takes the
  13. callback (provided by this class) as an argument
  14. """
  15. threading.Thread.__init__(self)
  16. self.function = function
  17. self.queue = queue
  18. self.die = False
  19. self.curl_hack = curl_hack
  20. def callback(self, data):
  21. try:
  22. if self.die:
  23. raise Exception() # trigger termination
  24. self.queue.put((1, data))
  25. except:
  26. if self.curl_hack:
  27. # We can't raise exceptions, because the pycurl
  28. # extension module will unconditionally print the
  29. # exception itself, and not pass it up to the caller.
  30. # Instead, just return a value that tells curl to
  31. # abort. (-1 would be best, in case we were given 0
  32. # bytes, but the extension doesn't support that).
  33. self.queue.put((2, sys.exc_info()))
  34. return 0
  35. raise
  36. def run(self):
  37. try:
  38. result = self.function(self.callback)
  39. except:
  40. self.queue.put((2, sys.exc_info()))
  41. else:
  42. self.queue.put((0, result))
  43. @contextlib.contextmanager
  44. def Iteratorizer(function, curl_hack = False):
  45. """
  46. Context manager that takes a function expecting a callback,
  47. and provides an iterable that yields the values passed to that
  48. callback instead.
  49. function: function to execute, which takes a callback
  50. (provided by this context manager) as an argument
  51. with iteratorizer(func) as it:
  52. for i in it:
  53. print 'callback was passed:', i
  54. print 'function returned:', it.retval
  55. """
  56. queue = Queue.Queue(maxsize = 1)
  57. thread = IteratorizerThread(queue, function, curl_hack)
  58. thread.daemon = True
  59. thread.start()
  60. class iteratorizer_gen(object):
  61. def __init__(self, queue):
  62. self.queue = queue
  63. self.retval = None
  64. def __iter__(self):
  65. return self
  66. def next(self):
  67. (typ, data) = self.queue.get()
  68. if typ == 0:
  69. # function has returned
  70. self.retval = data
  71. raise StopIteration
  72. elif typ == 1:
  73. # data is available
  74. return data
  75. else:
  76. # callback raised an exception
  77. raise data[0], data[1], data[2]
  78. try:
  79. yield iteratorizer_gen(queue)
  80. finally:
  81. # Ask the thread to die, if it's still running.
  82. thread.die = True
  83. while thread.isAlive():
  84. try:
  85. queue.get(True, 0.01)
  86. except:
  87. pass