1#!/usr/bin/env python
2# Copyright 2011 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Apply gzip/deflate to separate chunks of data."""
17
18import struct
19import zlib
20
21GZIP_HEADER = (
22    '\037\213'             # magic header
23    '\010'                 # compression method
24    '\000'                 # flags (none)
25    '\000\000\000\000'     # packed time (use zero)
26    '\002'
27    '\377')
28
29
30def compress_chunks(uncompressed_chunks, use_gzip):
31  """Compress a list of data with gzip or deflate.
32
33  The returned chunks may be used with HTTP chunked encoding.
34
35  Args:
36    uncompressed_chunks: a list of strings
37       (e.g. ["this is the first chunk", "and the second"])
38    use_gzip: if True, compress with gzip. Otherwise, use deflate.
39
40  Returns:
41    [compressed_chunk_1, compressed_chunk_2, ...]
42  """
43  if use_gzip:
44    size = 0
45    crc = zlib.crc32("") & 0xffffffffL
46    compressor = zlib.compressobj(
47        6, zlib.DEFLATED, -zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
48  else:
49    compressor = zlib.compressobj()
50  compressed_chunks = []
51  last_index = len(uncompressed_chunks) - 1
52  for index, data in enumerate(uncompressed_chunks):
53    chunk = ''
54    if use_gzip:
55      size += len(data)
56      crc = zlib.crc32(data, crc) & 0xffffffffL
57      if index == 0:
58        chunk += GZIP_HEADER
59    chunk += compressor.compress(data)
60    if index < last_index:
61      chunk += compressor.flush(zlib.Z_SYNC_FLUSH)
62    else:
63      chunk += (compressor.flush(zlib.Z_FULL_FLUSH) +
64                compressor.flush())
65      if use_gzip:
66        chunk += (struct.pack("<L", long(crc)) +
67                  struct.pack("<L", long(size)))
68    compressed_chunks.append(chunk)
69  return compressed_chunks
70
71
72def uncompress_chunks(compressed_chunks, use_gzip):
73  """Uncompress a list of data compressed with gzip or deflate.
74
75  Args:
76    compressed_chunks: a list of compressed data
77    use_gzip: if True, uncompress with gzip. Otherwise, use deflate.
78
79  Returns:
80    [uncompressed_chunk_1, uncompressed_chunk_2, ...]
81  """
82  if use_gzip:
83    decompress = zlib.decompressobj(16 + zlib.MAX_WBITS).decompress
84  else:
85    decompress = zlib.decompressobj(-zlib.MAX_WBITS).decompress
86  return [decompress(c) for c in compressed_chunks]
87