jpayne@68: # Copyright (C) 2001-2006 Python Software Foundation
jpayne@68: # Author: Barry Warsaw
jpayne@68: # Contact: email-sig@python.org
jpayne@68: 
jpayne@68: """Various types of useful iterators and generators."""
jpayne@68: 
jpayne@68: __all__ = [
jpayne@68:     'body_line_iterator',
jpayne@68:     'typed_subpart_iterator',
jpayne@68:     'walk',
jpayne@68:     # Do not include _structure() since it's part of the debugging API.
jpayne@68:     ]
jpayne@68: 
jpayne@68: import sys
jpayne@68: from io import StringIO
jpayne@68: 
jpayne@68: 
jpayne@68: 
jpayne@68: # This function will become a method of the Message class
jpayne@68: def walk(self):
jpayne@68:     """Walk over the message tree, yielding each subpart.
jpayne@68: 
jpayne@68:     The walk is performed in depth-first order.  This method is a
jpayne@68:     generator.
jpayne@68:     """
jpayne@68:     yield self
jpayne@68:     if self.is_multipart():
jpayne@68:         for subpart in self.get_payload():
jpayne@68:             yield from subpart.walk()
jpayne@68: 
jpayne@68: 
jpayne@68: 
jpayne@68: # These two functions are imported into the Iterators.py interface module.
jpayne@68: def body_line_iterator(msg, decode=False):
jpayne@68:     """Iterate over the parts, returning string payloads line-by-line.
jpayne@68: 
jpayne@68:     Optional decode (default False) is passed through to .get_payload().
jpayne@68:     """
jpayne@68:     for subpart in msg.walk():
jpayne@68:         payload = subpart.get_payload(decode=decode)
jpayne@68:         if isinstance(payload, str):
jpayne@68:             yield from StringIO(payload)
jpayne@68: 
jpayne@68: 
jpayne@68: def typed_subpart_iterator(msg, maintype='text', subtype=None):
jpayne@68:     """Iterate over the subparts with a given MIME type.
jpayne@68: 
jpayne@68:     Use `maintype' as the main MIME type to match against; this defaults to
jpayne@68:     "text".  Optional `subtype' is the MIME subtype to match against; if
jpayne@68:     omitted, only the main type is matched.
jpayne@68:     """
jpayne@68:     for subpart in msg.walk():
jpayne@68:         if subpart.get_content_maintype() == maintype:
jpayne@68:             if subtype is None or subpart.get_content_subtype() == subtype:
jpayne@68:                 yield subpart
jpayne@68: 
jpayne@68: 
jpayne@68: 
jpayne@68: def _structure(msg, fp=None, level=0, include_default=False):
jpayne@68:     """A handy debugging aid"""
jpayne@68:     if fp is None:
jpayne@68:         fp = sys.stdout
jpayne@68:     tab = ' ' * (level * 4)
jpayne@68:     print(tab + msg.get_content_type(), end='', file=fp)
jpayne@68:     if include_default:
jpayne@68:         print(' [%s]' % msg.get_default_type(), file=fp)
jpayne@68:     else:
jpayne@68:         print(file=fp)
jpayne@68:     if msg.is_multipart():
jpayne@68:         for subpart in msg.get_payload():
jpayne@68:             _structure(subpart, fp, level+1, include_default)