lib/testtools/testtools/content.py

   1 # Copyright (c) 2009-2011 testtools developers. See LICENSE for details.
   2
   3 """Content - a MIME-like Content object."""
   4
   5 __all__ = [
   6     'attach_file',
   7     'Content',
   8     'content_from_file',
   9     'content_from_stream',
  10     'text_content',
  11     'TracebackContent',
  12     ]
  13
  14 import codecs
  15 import os
  16
  17 from testtools import try_import
  18 from testtools.compat import _b
  19 from testtools.content_type import ContentType, UTF8_TEXT
  20 from testtools.testresult import TestResult
  21
  22 functools = try_import('functools')
  23
  24 _join_b = _b("").join
  25
  26
  27 DEFAULT_CHUNK_SIZE = 4096
  28
  29
  30 def _iter_chunks(stream, chunk_size):
  31     """Read 'stream' in chunks of 'chunk_size'.
  32
  33     :param stream: A file-like object to read from.
  34     :param chunk_size: The size of each read from 'stream'.
  35     """
  36     chunk = stream.read(chunk_size)
  37     while chunk:
  38         yield chunk
  39         chunk = stream.read(chunk_size)
  40
  41
  42 class Content(object):
  43     """A MIME-like Content object.
  44
  45     Content objects can be serialised to bytes using the iter_bytes method.
  46     If the Content-Type is recognised by other code, they are welcome to
  47     look for richer contents that mere byte serialisation - for example in
  48     memory object graphs etc. However, such code MUST be prepared to receive
  49     a generic Content object that has been reconstructed from a byte stream.
  50
  51     :ivar content_type: The content type of this Content.
  52     """
  53
  54     def __init__(self, content_type, get_bytes):
  55         """Create a ContentType."""
  56         if None in (content_type, get_bytes):
  57             raise ValueError("None not permitted in %r, %r" % (
  58                 content_type, get_bytes))
  59         self.content_type = content_type
  60         self._get_bytes = get_bytes
  61
  62     def __eq__(self, other):
  63         return (self.content_type == other.content_type and
  64             _join_b(self.iter_bytes()) == _join_b(other.iter_bytes()))
  65
  66     def iter_bytes(self):
  67         """Iterate over bytestrings of the serialised content."""
  68         return self._get_bytes()
  69
  70     def iter_text(self):
  71         """Iterate over the text of the serialised content.
  72
  73         This is only valid for text MIME types, and will use ISO-8859-1 if
  74         no charset parameter is present in the MIME type. (This is somewhat
  75         arbitrary, but consistent with RFC2617 3.7.1).
  76
  77         :raises ValueError: If the content type is not text/\*.
  78         """
  79         if self.content_type.type != "text":
  80             raise ValueError("Not a text type %r" % self.content_type)
  81         return self._iter_text()
  82
  83     def _iter_text(self):
  84         """Worker for iter_text - does the decoding."""
  85         encoding = self.content_type.parameters.get('charset', 'ISO-8859-1')
  86         try:
  87             # 2.5+
  88             decoder = codecs.getincrementaldecoder(encoding)()
  89             for bytes in self.iter_bytes():
  90                 yield decoder.decode(bytes)
  91             final = decoder.decode(_b(''), True)
  92             if final:
  93                 yield final
  94         except AttributeError:
  95             # < 2.5
  96             bytes = ''.join(self.iter_bytes())
  97             yield bytes.decode(encoding)
  98
  99     def __repr__(self):
 100         return "<Content type=%r, value=%r>" % (
 101             self.content_type, _join_b(self.iter_bytes()))
 102
 103
 104 class TracebackContent(Content):
 105     """Content object for tracebacks.
 106
 107     This adapts an exc_info tuple to the Content interface.
 108     text/x-traceback;language=python is used for the mime type, in order to
 109     provide room for other languages to format their tracebacks differently.
 110     """
 111
 112     def __init__(self, err, test):
 113         """Create a TracebackContent for err."""
 114         if err is None:
 115             raise ValueError("err may not be None")
 116         content_type = ContentType('text', 'x-traceback',
 117             {"language": "python", "charset": "utf8"})
 118         self._result = TestResult()
 119         value = self._result._exc_info_to_unicode(err, test)
 120         super(TracebackContent, self).__init__(
 121             content_type, lambda: [value.encode("utf8")])
 122
 123
 124 def text_content(text):
 125     """Create a `Content` object from some text.
 126
 127     This is useful for adding details which are short strings.
 128     """
 129     return Content(UTF8_TEXT, lambda: [text.encode('utf8')])
 130
 131
 132
 133 def maybe_wrap(wrapper, func):
 134     """Merge metadata for func into wrapper if functools is present."""
 135     if functools is not None:
 136         wrapper = functools.update_wrapper(wrapper, func)
 137     return wrapper
 138
 139
 140 def content_from_file(path, content_type=None, chunk_size=DEFAULT_CHUNK_SIZE,
 141                       buffer_now=False):
 142     """Create a `Content` object from a file on disk.
 143
 144     Note that unless 'read_now' is explicitly passed in as True, the file
 145     will only be read from when ``iter_bytes`` is called.
 146
 147     :param path: The path to the file to be used as content.
 148     :param content_type: The type of content.  If not specified, defaults
 149         to UTF8-encoded text/plain.
 150     :param chunk_size: The size of chunks to read from the file.
 151         Defaults to ``DEFAULT_CHUNK_SIZE``.
 152     :param buffer_now: If True, read the file from disk now and keep it in
 153         memory. Otherwise, only read when the content is serialized.
 154     """
 155     if content_type is None:
 156         content_type = UTF8_TEXT
 157     def reader():
 158         # This should be try:finally:, but python2.4 makes that hard. When
 159         # We drop older python support we can make this use a context manager
 160         # for maximum simplicity.
 161         stream = open(path, 'rb')
 162         for chunk in _iter_chunks(stream, chunk_size):
 163             yield chunk
 164         stream.close()
 165     return content_from_reader(reader, content_type, buffer_now)
 166
 167
 168 def content_from_stream(stream, content_type=None,
 169                         chunk_size=DEFAULT_CHUNK_SIZE, buffer_now=False):
 170     """Create a `Content` object from a file-like stream.
 171
 172     Note that the stream will only be read from when ``iter_bytes`` is
 173     called.
 174
 175     :param stream: A file-like object to read the content from. The stream
 176         is not closed by this function or the content object it returns.
 177     :param content_type: The type of content. If not specified, defaults
 178         to UTF8-encoded text/plain.
 179     :param chunk_size: The size of chunks to read from the file.
 180         Defaults to ``DEFAULT_CHUNK_SIZE``.
 181     :param buffer_now: If True, reads from the stream right now. Otherwise,
 182         only reads when the content is serialized. Defaults to False.
 183     """
 184     if content_type is None:
 185         content_type = UTF8_TEXT
 186     reader = lambda: _iter_chunks(stream, chunk_size)
 187     return content_from_reader(reader, content_type, buffer_now)
 188
 189
 190 def content_from_reader(reader, content_type, buffer_now):
 191     """Create a Content object that will obtain the content from reader.
 192
 193     :param reader: A callback to read the content. Should return an iterable of
 194         bytestrings.
 195     :param content_type: The content type to create.
 196     :param buffer_now: If True the reader is evaluated immediately and
 197         buffered.
 198     """
 199     if content_type is None:
 200         content_type = UTF8_TEXT
 201     if buffer_now:
 202         contents = list(reader())
 203         reader = lambda: contents
 204     return Content(content_type, reader)
 205
 206
 207 def attach_file(detailed, path, name=None, content_type=None,
 208                 chunk_size=DEFAULT_CHUNK_SIZE, buffer_now=True):
 209     """Attach a file to this test as a detail.
 210
 211     This is a convenience method wrapping around ``addDetail``.
 212
 213     Note that unless 'read_now' is explicitly passed in as True, the file
 214     *must* exist when the test result is called with the results of this
 215     test, after the test has been torn down.
 216
 217     :param detailed: An object with details
 218     :param path: The path to the file to attach.
 219     :param name: The name to give to the detail for the attached file.
 220     :param content_type: The content type of the file.  If not provided,
 221         defaults to UTF8-encoded text/plain.
 222     :param chunk_size: The size of chunks to read from the file.  Defaults
 223         to something sensible.
 224     :param buffer_now: If False the file content is read when the content
 225         object is evaluated rather than when attach_file is called.
 226         Note that this may be after any cleanups that obj_with_details has, so
 227         if the file is a temporary file disabling buffer_now may cause the file
 228         to be read after it is deleted. To handle those cases, using
 229         attach_file as a cleanup is recommended because it guarantees a
 230         sequence for when the attach_file call is made::
 231
 232             detailed.addCleanup(attach_file, 'foo.txt', detailed)
 233     """
 234     if name is None:
 235         name = os.path.basename(path)
 236     content_object = content_from_file(
 237         path, content_type, chunk_size, buffer_now)
 238     detailed.addDetail(name, content_object)