filewrapper.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. from io import BytesIO
  2. class CallbackFileWrapper(object):
  3. """
  4. Small wrapper around a fp object which will tee everything read into a
  5. buffer, and when that file is closed it will execute a callback with the
  6. contents of that buffer.
  7. All attributes are proxied to the underlying file object.
  8. This class uses members with a double underscore (__) leading prefix so as
  9. not to accidentally shadow an attribute.
  10. """
  11. def __init__(self, fp, callback):
  12. self.__buf = BytesIO()
  13. self.__fp = fp
  14. self.__callback = callback
  15. def __getattr__(self, name):
  16. # The vaguaries of garbage collection means that self.__fp is
  17. # not always set. By using __getattribute__ and the private
  18. # name[0] allows looking up the attribute value and raising an
  19. # AttributeError when it doesn't exist. This stop thigns from
  20. # infinitely recursing calls to getattr in the case where
  21. # self.__fp hasn't been set.
  22. #
  23. # [0] https://docs.python.org/2/reference/expressions.html#atom-identifiers
  24. fp = self.__getattribute__('_CallbackFileWrapper__fp')
  25. return getattr(fp, name)
  26. def __is_fp_closed(self):
  27. try:
  28. return self.__fp.fp is None
  29. except AttributeError:
  30. pass
  31. try:
  32. return self.__fp.closed
  33. except AttributeError:
  34. pass
  35. # We just don't cache it then.
  36. # TODO: Add some logging here...
  37. return False
  38. def _close(self):
  39. if self.__callback:
  40. self.__callback(self.__buf.getvalue())
  41. # We assign this to None here, because otherwise we can get into
  42. # really tricky problems where the CPython interpreter dead locks
  43. # because the callback is holding a reference to something which
  44. # has a __del__ method. Setting this to None breaks the cycle
  45. # and allows the garbage collector to do it's thing normally.
  46. self.__callback = None
  47. def read(self, amt=None):
  48. data = self.__fp.read(amt)
  49. self.__buf.write(data)
  50. if self.__is_fp_closed():
  51. self._close()
  52. return data
  53. def _safe_read(self, amt):
  54. data = self.__fp._safe_read(amt)
  55. if amt == 2 and data == b'\r\n':
  56. # urllib executes this read to toss the CRLF at the end
  57. # of the chunk.
  58. return data
  59. self.__buf.write(data)
  60. if self.__is_fp_closed():
  61. self._close()
  62. return data