iterators.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. # Copyright (C) 2001-2006 Python Software Foundation
  2. # Author: Barry Warsaw
  3. # Contact: email-sig@python.org
  4. """Various types of useful iterators and generators."""
  5. from __future__ import print_function
  6. from __future__ import unicode_literals
  7. from __future__ import division
  8. from __future__ import absolute_import
  9. __all__ = [
  10. 'body_line_iterator',
  11. 'typed_subpart_iterator',
  12. 'walk',
  13. # Do not include _structure() since it's part of the debugging API.
  14. ]
  15. import sys
  16. from io import StringIO
  17. # This function will become a method of the Message class
  18. def walk(self):
  19. """Walk over the message tree, yielding each subpart.
  20. The walk is performed in depth-first order. This method is a
  21. generator.
  22. """
  23. yield self
  24. if self.is_multipart():
  25. for subpart in self.get_payload():
  26. for subsubpart in subpart.walk():
  27. yield subsubpart
  28. # These two functions are imported into the Iterators.py interface module.
  29. def body_line_iterator(msg, decode=False):
  30. """Iterate over the parts, returning string payloads line-by-line.
  31. Optional decode (default False) is passed through to .get_payload().
  32. """
  33. for subpart in msg.walk():
  34. payload = subpart.get_payload(decode=decode)
  35. if isinstance(payload, str):
  36. for line in StringIO(payload):
  37. yield line
  38. def typed_subpart_iterator(msg, maintype='text', subtype=None):
  39. """Iterate over the subparts with a given MIME type.
  40. Use `maintype' as the main MIME type to match against; this defaults to
  41. "text". Optional `subtype' is the MIME subtype to match against; if
  42. omitted, only the main type is matched.
  43. """
  44. for subpart in msg.walk():
  45. if subpart.get_content_maintype() == maintype:
  46. if subtype is None or subpart.get_content_subtype() == subtype:
  47. yield subpart
  48. def _structure(msg, fp=None, level=0, include_default=False):
  49. """A handy debugging aid"""
  50. if fp is None:
  51. fp = sys.stdout
  52. tab = ' ' * (level * 4)
  53. print(tab + msg.get_content_type(), end='', file=fp)
  54. if include_default:
  55. print(' [%s]' % msg.get_default_type(), file=fp)
  56. else:
  57. print(file=fp)
  58. if msg.is_multipart():
  59. for subpart in msg.get_payload():
  60. _structure(subpart, fp, level+1, include_default)