Guido van Rossum | 8b3febe | 2007-08-30 01:15:14 +0000 | [diff] [blame] | 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 | |
| 19 | # This function will become a method of the Message class |
| 20 | def walk(self): |
| 21 | """Walk over the message tree, yielding each subpart. |
| 22 | |
| 23 | The walk is performed in depth-first order. This method is a |
| 24 | generator. |
| 25 | """ |
| 26 | yield self |
| 27 | if self.is_multipart(): |
| 28 | for subpart in self.get_payload(): |
Philip Jenvey | 4993cc0 | 2012-10-01 12:53:43 -0700 | [diff] [blame] | 29 | yield from subpart.walk() |
Guido van Rossum | 8b3febe | 2007-08-30 01:15:14 +0000 | [diff] [blame] | 30 | |
| 31 | |
| 32 | |
| 33 | # These two functions are imported into the Iterators.py interface module. |
| 34 | def body_line_iterator(msg, decode=False): |
| 35 | """Iterate over the parts, returning string payloads line-by-line. |
| 36 | |
| 37 | Optional decode (default False) is passed through to .get_payload(). |
| 38 | """ |
| 39 | for subpart in msg.walk(): |
| 40 | payload = subpart.get_payload(decode=decode) |
Guido van Rossum | 3172c5d | 2007-10-16 18:12:55 +0000 | [diff] [blame] | 41 | if isinstance(payload, str): |
Philip Jenvey | 4993cc0 | 2012-10-01 12:53:43 -0700 | [diff] [blame] | 42 | yield from StringIO(payload) |
Guido van Rossum | 8b3febe | 2007-08-30 01:15:14 +0000 | [diff] [blame] | 43 | |
| 44 | |
| 45 | def typed_subpart_iterator(msg, maintype='text', subtype=None): |
| 46 | """Iterate over the subparts with a given MIME type. |
| 47 | |
| 48 | Use `maintype' as the main MIME type to match against; this defaults to |
| 49 | "text". Optional `subtype' is the MIME subtype to match against; if |
| 50 | omitted, only the main type is matched. |
| 51 | """ |
| 52 | for subpart in msg.walk(): |
| 53 | if subpart.get_content_maintype() == maintype: |
| 54 | if subtype is None or subpart.get_content_subtype() == subtype: |
| 55 | yield subpart |
| 56 | |
| 57 | |
| 58 | |
| 59 | def _structure(msg, fp=None, level=0, include_default=False): |
| 60 | """A handy debugging aid""" |
| 61 | if fp is None: |
| 62 | fp = sys.stdout |
| 63 | tab = ' ' * (level * 4) |
| 64 | print(tab + msg.get_content_type(), end='', file=fp) |
| 65 | if include_default: |
| 66 | print(' [%s]' % msg.get_default_type(), file=fp) |
| 67 | else: |
| 68 | print(file=fp) |
| 69 | if msg.is_multipart(): |
| 70 | for subpart in msg.get_payload(): |
| 71 | _structure(subpart, fp, level+1, include_default) |