iterators.py 2.1 KB

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