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