summaryrefslogtreecommitdiff
path: root/lib/fixtures/_fixtures/streams.py
blob: 654ee67bf8a1b248b861a0c514a481c422534905 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#  fixtures: Fixtures with cleanups for testing and convenience.
#
# Copyright (c) 2012, Robert Collins <robertc@robertcollins.net>
#
# Licensed under either the Apache License, Version 2.0 or the BSD 3-clause
# license at the users choice. A copy of both licenses are available in the
# project source as Apache-2.0 and BSD. You may not use this file except in
# compliance with one of these two licences.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under these licenses is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# license you chose for the specific language governing permissions and
# limitations under that license.

__all__ = [
    'ByteStream',
    'DetailStream',
    'StringStream',
    ]

import io
import sys

from fixtures import Fixture
import testtools


class Stream(Fixture):
    """Expose a file-like object as a detail.

    :attr stream: The file-like object.
    """

    def __init__(self, detail_name, stream_factory):
        """Create a ByteStream.

        :param detail_name: Use this as the name of the stream.
        :param stream_factory: Called to construct a pair of streams:
            (write_stream, content_stream).
        """
        self._detail_name = detail_name
        self._stream_factory = stream_factory

    def setUp(self):
        super(Stream, self).setUp()
        write_stream, read_stream = self._stream_factory()
        self.stream = write_stream
        self.addDetail(self._detail_name,
            testtools.content.content_from_stream(read_stream, seek_offset=0))


def _byte_stream_factory():
    result = io.BytesIO()
    return (result, result)


def ByteStream(detail_name):
    """Provide a file-like object that accepts bytes and expose as a detail.

    :param detail_name: The name of the detail.
    :return: A fixture which has an attribute `stream` containing the file-like
        object.
    """
    return Stream(detail_name, _byte_stream_factory)


def _string_stream_factory():
    lower = io.BytesIO()
    upper = io.TextIOWrapper(lower, encoding="utf8")
    # In theory, this is sufficient and correct, but on Python2,
    # upper.write(_b('foo")) will whinge louadly.
    if sys.version_info[0] < 3:
        upper_write = upper.write
        def safe_write(str_or_bytes):
            if type(str_or_bytes) is str:
                str_or_bytes = str_or_bytes.decode('utf8')
            return upper_write(str_or_bytes)
        upper.write = safe_write
    return upper, lower


def StringStream(detail_name):
    """Provide a file-like object that accepts strings and expose as a detail.

    :param detail_name: The name of the detail.
    :return: A fixture which has an attribute `stream` containing the file-like
        object.
    """
    return Stream(detail_name, _string_stream_factory)


def DetailStream(detail_name):
    """Deprecated alias for ByteStream."""
    return ByteStream(detail_name)