diff options
Diffstat (limited to 'lib/subunit/python/subunit/chunked.py')
-rw-r--r-- | lib/subunit/python/subunit/chunked.py | 164 |
1 files changed, 164 insertions, 0 deletions
diff --git a/lib/subunit/python/subunit/chunked.py b/lib/subunit/python/subunit/chunked.py new file mode 100644 index 0000000000..82e4b0ddfc --- /dev/null +++ b/lib/subunit/python/subunit/chunked.py @@ -0,0 +1,164 @@ +# +# subunit: extensions to python unittest to get test results from subprocesses. +# Copyright (C) 2005 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. +# + +"""Encoder/decoder for http style chunked encoding.""" + +class Decoder(object): + """Decode chunked content to a byte stream.""" + + def __init__(self, output): + """Create a decoder decoding to output. + + :param output: A file-like object. Bytes written to the Decoder are + decoded to strip off the chunking and written to the output. + Up to a full write worth of data or a single control line may be + buffered (whichever is larger). The close method should be called + when no more data is available, to detect short streams; the + write method will return none-None when the end of a stream is + detected. + """ + self.output = output + self.buffered_bytes = [] + self.state = self._read_length + self.body_length = 0 + + def close(self): + """Close the decoder. + + :raises ValueError: If the stream is incomplete ValueError is raised. + """ + if self.state != self._finished: + raise ValueError("incomplete stream") + + def _finished(self): + """Finished reading, return any remaining bytes.""" + if self.buffered_bytes: + buffered_bytes = self.buffered_bytes + self.buffered_bytes = [] + return ''.join(buffered_bytes) + else: + raise ValueError("stream is finished") + + def _read_body(self): + """Pass body bytes to the output.""" + while self.body_length and self.buffered_bytes: + if self.body_length >= len(self.buffered_bytes[0]): + self.output.write(self.buffered_bytes[0]) + self.body_length -= len(self.buffered_bytes[0]) + del self.buffered_bytes[0] + # No more data available. + if not self.body_length: + self.state = self._read_length + else: + self.output.write(self.buffered_bytes[0][:self.body_length]) + self.buffered_bytes[0] = \ + self.buffered_bytes[0][self.body_length:] + self.body_length = 0 + self.state = self._read_length + return self.state() + + def _read_length(self): + """Try to decode a length from the bytes.""" + count = -1 + match_chars = "0123456789abcdefABCDEF\r\n" + count_chars = [] + for bytes in self.buffered_bytes: + for byte in bytes: + if byte not in match_chars: + break + count_chars.append(byte) + if byte == '\n': + break + if not count_chars: + return + if count_chars[-1][-1] != '\n': + return + count_str = ''.join(count_chars) + self.body_length = int(count_str[:-2], 16) + excess_bytes = len(count_str) + while excess_bytes: + if excess_bytes >= len(self.buffered_bytes[0]): + excess_bytes -= len(self.buffered_bytes[0]) + del self.buffered_bytes[0] + else: + self.buffered_bytes[0] = self.buffered_bytes[0][excess_bytes:] + excess_bytes = 0 + if not self.body_length: + self.state = self._finished + if not self.buffered_bytes: + # May not call into self._finished with no buffered data. + return '' + else: + self.state = self._read_body + return self.state() + + def write(self, bytes): + """Decode bytes to the output stream. + + :raises ValueError: If the stream has already seen the end of file + marker. + :returns: None, or the excess bytes beyond the end of file marker. + """ + if bytes: + self.buffered_bytes.append(bytes) + return self.state() + + +class Encoder(object): + """Encode content to a stream using HTTP Chunked coding.""" + + def __init__(self, output): + """Create an encoder encoding to output. + + :param output: A file-like object. Bytes written to the Encoder + will be encoded using HTTP chunking. Small writes may be buffered + and the ``close`` method must be called to finish the stream. + """ + self.output = output + self.buffered_bytes = [] + self.buffer_size = 0 + + def flush(self, extra_len=0): + """Flush the encoder to the output stream. + + :param extra_len: Increase the size of the chunk by this many bytes + to allow for a subsequent write. + """ + if not self.buffer_size and not extra_len: + return + buffered_bytes = self.buffered_bytes + buffer_size = self.buffer_size + self.buffered_bytes = [] + self.buffer_size = 0 + self.output.write("%X\r\n" % (buffer_size + extra_len)) + if buffer_size: + self.output.write(''.join(buffered_bytes)) + return True + + def write(self, bytes): + """Encode bytes to the output stream.""" + bytes_len = len(bytes) + if self.buffer_size + bytes_len >= 65536: + self.flush(bytes_len) + self.output.write(bytes) + else: + self.buffered_bytes.append(bytes) + self.buffer_size += bytes_len + + def close(self): + """Finish the stream. This does not close the output stream.""" + self.flush() + self.output.write("0\r\n") |