126fa459cSmrg#define PY_SSIZE_T_CLEAN 1
226fa459cSmrg#include <Python.h>
326fa459cSmrg#include <bytesobject.h>
426fa459cSmrg#include <structmember.h>
526fa459cSmrg#include <vector>
626fa459cSmrg#include "../common/version.h"
726fa459cSmrg#include <brotli/decode.h>
826fa459cSmrg#include <brotli/encode.h>
926fa459cSmrg
1026fa459cSmrg#if PY_MAJOR_VERSION >= 3
1126fa459cSmrg#define PyInt_Check PyLong_Check
1226fa459cSmrg#define PyInt_AsLong PyLong_AsLong
1326fa459cSmrg#endif
1426fa459cSmrg
1526fa459cSmrgstatic PyObject *BrotliError;
1626fa459cSmrg
1726fa459cSmrgstatic int as_bounded_int(PyObject *o, int* result, int lower_bound, int upper_bound) {
1826fa459cSmrg  long value = PyInt_AsLong(o);
1926fa459cSmrg  if ((value < (long) lower_bound) || (value > (long) upper_bound)) {
2026fa459cSmrg    return 0;
2126fa459cSmrg  }
2226fa459cSmrg  *result = (int) value;
2326fa459cSmrg  return 1;
2426fa459cSmrg}
2526fa459cSmrg
2626fa459cSmrgstatic int mode_convertor(PyObject *o, BrotliEncoderMode *mode) {
2726fa459cSmrg  if (!PyInt_Check(o)) {
2826fa459cSmrg    PyErr_SetString(BrotliError, "Invalid mode");
2926fa459cSmrg    return 0;
3026fa459cSmrg  }
3126fa459cSmrg
3226fa459cSmrg  int mode_value = -1;
3326fa459cSmrg  if (!as_bounded_int(o, &mode_value, 0, 255)) {
3426fa459cSmrg    PyErr_SetString(BrotliError, "Invalid mode");
3526fa459cSmrg    return 0;
3626fa459cSmrg  }
3726fa459cSmrg  *mode = (BrotliEncoderMode) mode_value;
3826fa459cSmrg  if (*mode != BROTLI_MODE_GENERIC &&
3926fa459cSmrg      *mode != BROTLI_MODE_TEXT &&
4026fa459cSmrg      *mode != BROTLI_MODE_FONT) {
4126fa459cSmrg    PyErr_SetString(BrotliError, "Invalid mode");
4226fa459cSmrg    return 0;
4326fa459cSmrg  }
4426fa459cSmrg
4526fa459cSmrg  return 1;
4626fa459cSmrg}
4726fa459cSmrg
4826fa459cSmrgstatic int quality_convertor(PyObject *o, int *quality) {
4926fa459cSmrg  if (!PyInt_Check(o)) {
5026fa459cSmrg    PyErr_SetString(BrotliError, "Invalid quality");
5126fa459cSmrg    return 0;
5226fa459cSmrg  }
5326fa459cSmrg
5426fa459cSmrg  if (!as_bounded_int(o, quality, 0, 11)) {
5526fa459cSmrg    PyErr_SetString(BrotliError, "Invalid quality. Range is 0 to 11.");
5626fa459cSmrg    return 0;
5726fa459cSmrg  }
5826fa459cSmrg
5926fa459cSmrg  return 1;
6026fa459cSmrg}
6126fa459cSmrg
6226fa459cSmrgstatic int lgwin_convertor(PyObject *o, int *lgwin) {
6326fa459cSmrg  if (!PyInt_Check(o)) {
6426fa459cSmrg    PyErr_SetString(BrotliError, "Invalid lgwin");
6526fa459cSmrg    return 0;
6626fa459cSmrg  }
6726fa459cSmrg
6826fa459cSmrg  if (!as_bounded_int(o, lgwin, 10, 24)) {
6926fa459cSmrg    PyErr_SetString(BrotliError, "Invalid lgwin. Range is 10 to 24.");
7026fa459cSmrg    return 0;
7126fa459cSmrg  }
7226fa459cSmrg
7326fa459cSmrg  return 1;
7426fa459cSmrg}
7526fa459cSmrg
7626fa459cSmrgstatic int lgblock_convertor(PyObject *o, int *lgblock) {
7726fa459cSmrg  if (!PyInt_Check(o)) {
7826fa459cSmrg    PyErr_SetString(BrotliError, "Invalid lgblock");
7926fa459cSmrg    return 0;
8026fa459cSmrg  }
8126fa459cSmrg
8226fa459cSmrg  if (!as_bounded_int(o, lgblock, 0, 24) || (*lgblock != 0 && *lgblock < 16)) {
8326fa459cSmrg    PyErr_SetString(BrotliError, "Invalid lgblock. Can be 0 or in range 16 to 24.");
8426fa459cSmrg    return 0;
8526fa459cSmrg  }
8626fa459cSmrg
8726fa459cSmrg  return 1;
8826fa459cSmrg}
8926fa459cSmrg
9026fa459cSmrgstatic BROTLI_BOOL compress_stream(BrotliEncoderState* enc, BrotliEncoderOperation op,
9126fa459cSmrg                                   std::vector<uint8_t>* output,
9226fa459cSmrg                                   uint8_t* input, size_t input_length) {
9326fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
9426fa459cSmrg  Py_BEGIN_ALLOW_THREADS
9526fa459cSmrg
9626fa459cSmrg  size_t available_in = input_length;
9726fa459cSmrg  const uint8_t* next_in = input;
9826fa459cSmrg  size_t available_out = 0;
9926fa459cSmrg  uint8_t* next_out = NULL;
10026fa459cSmrg
10126fa459cSmrg  while (ok) {
10226fa459cSmrg    ok = BrotliEncoderCompressStream(enc, op,
10326fa459cSmrg                                     &available_in, &next_in,
10426fa459cSmrg                                     &available_out, &next_out, NULL);
10526fa459cSmrg    if (!ok)
10626fa459cSmrg      break;
10726fa459cSmrg
10826fa459cSmrg    size_t buffer_length = 0; // Request all available output.
10926fa459cSmrg    const uint8_t* buffer = BrotliEncoderTakeOutput(enc, &buffer_length);
11026fa459cSmrg    if (buffer_length) {
11126fa459cSmrg      (*output).insert((*output).end(), buffer, buffer + buffer_length);
11226fa459cSmrg    }
11326fa459cSmrg
11426fa459cSmrg    if (available_in || BrotliEncoderHasMoreOutput(enc)) {
11526fa459cSmrg      continue;
11626fa459cSmrg    }
11726fa459cSmrg
11826fa459cSmrg    break;
11926fa459cSmrg  }
12026fa459cSmrg
12126fa459cSmrg  Py_END_ALLOW_THREADS
12226fa459cSmrg  return ok;
12326fa459cSmrg}
12426fa459cSmrg
12526fa459cSmrgPyDoc_STRVAR(brotli_Compressor_doc,
12626fa459cSmrg"An object to compress a byte string.\n"
12726fa459cSmrg"\n"
12826fa459cSmrg"Signature:\n"
12926fa459cSmrg"  Compressor(mode=MODE_GENERIC, quality=11, lgwin=22, lgblock=0)\n"
13026fa459cSmrg"\n"
13126fa459cSmrg"Args:\n"
13226fa459cSmrg"  mode (int, optional): The compression mode can be MODE_GENERIC (default),\n"
13326fa459cSmrg"    MODE_TEXT (for UTF-8 format text input) or MODE_FONT (for WOFF 2.0). \n"
13426fa459cSmrg"  quality (int, optional): Controls the compression-speed vs compression-\n"
13526fa459cSmrg"    density tradeoff. The higher the quality, the slower the compression.\n"
13626fa459cSmrg"    Range is 0 to 11. Defaults to 11.\n"
13726fa459cSmrg"  lgwin (int, optional): Base 2 logarithm of the sliding window size. Range\n"
13826fa459cSmrg"    is 10 to 24. Defaults to 22.\n"
13926fa459cSmrg"  lgblock (int, optional): Base 2 logarithm of the maximum input block size.\n"
14026fa459cSmrg"    Range is 16 to 24. If set to 0, the value will be set based on the\n"
14126fa459cSmrg"    quality. Defaults to 0.\n"
14226fa459cSmrg"\n"
14326fa459cSmrg"Raises:\n"
14426fa459cSmrg"  brotli.error: If arguments are invalid.\n");
14526fa459cSmrg
14626fa459cSmrgtypedef struct {
14726fa459cSmrg  PyObject_HEAD
14826fa459cSmrg  BrotliEncoderState* enc;
14926fa459cSmrg} brotli_Compressor;
15026fa459cSmrg
15126fa459cSmrgstatic void brotli_Compressor_dealloc(brotli_Compressor* self) {
15226fa459cSmrg  BrotliEncoderDestroyInstance(self->enc);
15326fa459cSmrg  #if PY_MAJOR_VERSION >= 3
15426fa459cSmrg  Py_TYPE(self)->tp_free((PyObject*)self);
15526fa459cSmrg  #else
15626fa459cSmrg  self->ob_type->tp_free((PyObject*)self);
15726fa459cSmrg  #endif
15826fa459cSmrg}
15926fa459cSmrg
16026fa459cSmrgstatic PyObject* brotli_Compressor_new(PyTypeObject *type, PyObject *args, PyObject *keywds) {
16126fa459cSmrg  brotli_Compressor *self;
16226fa459cSmrg  self = (brotli_Compressor *)type->tp_alloc(type, 0);
16326fa459cSmrg
16426fa459cSmrg  if (self != NULL) {
16526fa459cSmrg    self->enc = BrotliEncoderCreateInstance(0, 0, 0);
16626fa459cSmrg  }
16726fa459cSmrg
16826fa459cSmrg  return (PyObject *)self;
16926fa459cSmrg}
17026fa459cSmrg
17126fa459cSmrgstatic int brotli_Compressor_init(brotli_Compressor *self, PyObject *args, PyObject *keywds) {
17226fa459cSmrg  BrotliEncoderMode mode = (BrotliEncoderMode) -1;
17326fa459cSmrg  int quality = -1;
17426fa459cSmrg  int lgwin = -1;
17526fa459cSmrg  int lgblock = -1;
17626fa459cSmrg  int ok;
17726fa459cSmrg
17826fa459cSmrg  static const char *kwlist[] = {"mode", "quality", "lgwin", "lgblock", NULL};
17926fa459cSmrg
18026fa459cSmrg  ok = PyArg_ParseTupleAndKeywords(args, keywds, "|O&O&O&O&:Compressor",
18126fa459cSmrg                    const_cast<char **>(kwlist),
18226fa459cSmrg                    &mode_convertor, &mode,
18326fa459cSmrg                    &quality_convertor, &quality,
18426fa459cSmrg                    &lgwin_convertor, &lgwin,
18526fa459cSmrg                    &lgblock_convertor, &lgblock);
18626fa459cSmrg  if (!ok)
18726fa459cSmrg    return -1;
18826fa459cSmrg  if (!self->enc)
18926fa459cSmrg    return -1;
19026fa459cSmrg
19126fa459cSmrg  if ((int) mode != -1)
19226fa459cSmrg    BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_MODE, (uint32_t)mode);
19326fa459cSmrg  if (quality != -1)
19426fa459cSmrg    BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_QUALITY, (uint32_t)quality);
19526fa459cSmrg  if (lgwin != -1)
19626fa459cSmrg    BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_LGWIN, (uint32_t)lgwin);
19726fa459cSmrg  if (lgblock != -1)
19826fa459cSmrg    BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_LGBLOCK, (uint32_t)lgblock);
19926fa459cSmrg
20026fa459cSmrg  return 0;
20126fa459cSmrg}
20226fa459cSmrg
20326fa459cSmrgPyDoc_STRVAR(brotli_Compressor_process_doc,
20426fa459cSmrg"Process \"string\" for compression, returning a string that contains \n"
20526fa459cSmrg"compressed output data.  This data should be concatenated to the output \n"
20626fa459cSmrg"produced by any preceding calls to the \"process()\" or flush()\" methods. \n"
20726fa459cSmrg"Some or all of the input may be kept in internal buffers for later \n"
20826fa459cSmrg"processing, and the compressed output data may be empty until enough input \n"
20926fa459cSmrg"has been accumulated.\n"
21026fa459cSmrg"\n"
21126fa459cSmrg"Signature:\n"
21226fa459cSmrg"  compress(string)\n"
21326fa459cSmrg"\n"
21426fa459cSmrg"Args:\n"
21526fa459cSmrg"  string (bytes): The input data\n"
21626fa459cSmrg"\n"
21726fa459cSmrg"Returns:\n"
21826fa459cSmrg"  The compressed output data (bytes)\n"
21926fa459cSmrg"\n"
22026fa459cSmrg"Raises:\n"
22126fa459cSmrg"  brotli.error: If compression fails\n");
22226fa459cSmrg
22326fa459cSmrgstatic PyObject* brotli_Compressor_process(brotli_Compressor *self, PyObject *args) {
22426fa459cSmrg  PyObject* ret = NULL;
22526fa459cSmrg  std::vector<uint8_t> output;
22626fa459cSmrg  Py_buffer input;
22726fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
22826fa459cSmrg
22926fa459cSmrg#if PY_MAJOR_VERSION >= 3
23026fa459cSmrg  ok = (BROTLI_BOOL)PyArg_ParseTuple(args, "y*:process", &input);
23126fa459cSmrg#else
23226fa459cSmrg  ok = (BROTLI_BOOL)PyArg_ParseTuple(args, "s*:process", &input);
23326fa459cSmrg#endif
23426fa459cSmrg
23526fa459cSmrg  if (!ok)
23626fa459cSmrg    return NULL;
23726fa459cSmrg
23826fa459cSmrg  if (!self->enc) {
23926fa459cSmrg    ok = BROTLI_FALSE;
24026fa459cSmrg    goto end;
24126fa459cSmrg  }
24226fa459cSmrg
24326fa459cSmrg  ok = compress_stream(self->enc, BROTLI_OPERATION_PROCESS,
24426fa459cSmrg                       &output, static_cast<uint8_t*>(input.buf), input.len);
24526fa459cSmrg
24626fa459cSmrgend:
24726fa459cSmrg  PyBuffer_Release(&input);
24826fa459cSmrg  if (ok) {
24926fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.size() ? &output[0] : NULL), output.size());
25026fa459cSmrg  } else {
25126fa459cSmrg    PyErr_SetString(BrotliError, "BrotliEncoderCompressStream failed while processing the stream");
25226fa459cSmrg  }
25326fa459cSmrg
25426fa459cSmrg  return ret;
25526fa459cSmrg}
25626fa459cSmrg
25726fa459cSmrgPyDoc_STRVAR(brotli_Compressor_flush_doc,
25826fa459cSmrg"Process all pending input, returning a string containing the remaining\n"
25926fa459cSmrg"compressed data. This data should be concatenated to the output produced by\n"
26026fa459cSmrg"any preceding calls to the \"process()\" or \"flush()\" methods.\n"
26126fa459cSmrg"\n"
26226fa459cSmrg"Signature:\n"
26326fa459cSmrg"  flush()\n"
26426fa459cSmrg"\n"
26526fa459cSmrg"Returns:\n"
26626fa459cSmrg"  The compressed output data (bytes)\n"
26726fa459cSmrg"\n"
26826fa459cSmrg"Raises:\n"
26926fa459cSmrg"  brotli.error: If compression fails\n");
27026fa459cSmrg
27126fa459cSmrgstatic PyObject* brotli_Compressor_flush(brotli_Compressor *self) {
27226fa459cSmrg  PyObject *ret = NULL;
27326fa459cSmrg  std::vector<uint8_t> output;
27426fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
27526fa459cSmrg
27626fa459cSmrg  if (!self->enc) {
27726fa459cSmrg    ok = BROTLI_FALSE;
27826fa459cSmrg    goto end;
27926fa459cSmrg  }
28026fa459cSmrg
28126fa459cSmrg  ok = compress_stream(self->enc, BROTLI_OPERATION_FLUSH,
28226fa459cSmrg                       &output, NULL, 0);
28326fa459cSmrg
28426fa459cSmrgend:
28526fa459cSmrg  if (ok) {
28626fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.size() ? &output[0] : NULL), output.size());
28726fa459cSmrg  } else {
28826fa459cSmrg    PyErr_SetString(BrotliError, "BrotliEncoderCompressStream failed while flushing the stream");
28926fa459cSmrg  }
29026fa459cSmrg
29126fa459cSmrg  return ret;
29226fa459cSmrg}
29326fa459cSmrg
29426fa459cSmrgPyDoc_STRVAR(brotli_Compressor_finish_doc,
29526fa459cSmrg"Process all pending input and complete all compression, returning a string\n"
29626fa459cSmrg"containing the remaining compressed data. This data should be concatenated\n"
29726fa459cSmrg"to the output produced by any preceding calls to the \"process()\" or\n"
29826fa459cSmrg"\"flush()\" methods.\n"
29926fa459cSmrg"After calling \"finish()\", the \"process()\" and \"flush()\" methods\n"
30026fa459cSmrg"cannot be called again, and a new \"Compressor\" object should be created.\n"
30126fa459cSmrg"\n"
30226fa459cSmrg"Signature:\n"
30326fa459cSmrg"  finish(string)\n"
30426fa459cSmrg"\n"
30526fa459cSmrg"Returns:\n"
30626fa459cSmrg"  The compressed output data (bytes)\n"
30726fa459cSmrg"\n"
30826fa459cSmrg"Raises:\n"
30926fa459cSmrg"  brotli.error: If compression fails\n");
31026fa459cSmrg
31126fa459cSmrgstatic PyObject* brotli_Compressor_finish(brotli_Compressor *self) {
31226fa459cSmrg  PyObject *ret = NULL;
31326fa459cSmrg  std::vector<uint8_t> output;
31426fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
31526fa459cSmrg
31626fa459cSmrg  if (!self->enc) {
31726fa459cSmrg    ok = BROTLI_FALSE;
31826fa459cSmrg    goto end;
31926fa459cSmrg  }
32026fa459cSmrg
32126fa459cSmrg  ok = compress_stream(self->enc, BROTLI_OPERATION_FINISH,
32226fa459cSmrg                       &output, NULL, 0);
32326fa459cSmrg
32426fa459cSmrg  if (ok) {
32526fa459cSmrg    ok = BrotliEncoderIsFinished(self->enc);
32626fa459cSmrg  }
32726fa459cSmrg
32826fa459cSmrgend:
32926fa459cSmrg  if (ok) {
33026fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.empty() ? NULL : &output[0]), output.size());
33126fa459cSmrg  } else {
33226fa459cSmrg    PyErr_SetString(BrotliError, "BrotliEncoderCompressStream failed while finishing the stream");
33326fa459cSmrg  }
33426fa459cSmrg
33526fa459cSmrg  return ret;
33626fa459cSmrg}
33726fa459cSmrg
33826fa459cSmrgstatic PyMemberDef brotli_Compressor_members[] = {
33926fa459cSmrg  {NULL}  /* Sentinel */
34026fa459cSmrg};
34126fa459cSmrg
34226fa459cSmrgstatic PyMethodDef brotli_Compressor_methods[] = {
34326fa459cSmrg  {"process", (PyCFunction)brotli_Compressor_process, METH_VARARGS, brotli_Compressor_process_doc},
34426fa459cSmrg  {"flush", (PyCFunction)brotli_Compressor_flush, METH_NOARGS, brotli_Compressor_flush_doc},
34526fa459cSmrg  {"finish", (PyCFunction)brotli_Compressor_finish, METH_NOARGS, brotli_Compressor_finish_doc},
34626fa459cSmrg  {NULL}  /* Sentinel */
34726fa459cSmrg};
34826fa459cSmrg
34926fa459cSmrgstatic PyTypeObject brotli_CompressorType = {
35026fa459cSmrg  #if PY_MAJOR_VERSION >= 3
35126fa459cSmrg  PyVarObject_HEAD_INIT(NULL, 0)
35226fa459cSmrg  #else
35326fa459cSmrg  PyObject_HEAD_INIT(NULL)
35426fa459cSmrg  0,                                     /* ob_size*/
35526fa459cSmrg  #endif
35626fa459cSmrg  "brotli.Compressor",                   /* tp_name */
35726fa459cSmrg  sizeof(brotli_Compressor),             /* tp_basicsize */
35826fa459cSmrg  0,                                     /* tp_itemsize */
35926fa459cSmrg  (destructor)brotli_Compressor_dealloc, /* tp_dealloc */
36026fa459cSmrg  0,                                     /* tp_print */
36126fa459cSmrg  0,                                     /* tp_getattr */
36226fa459cSmrg  0,                                     /* tp_setattr */
36326fa459cSmrg  0,                                     /* tp_compare */
36426fa459cSmrg  0,                                     /* tp_repr */
36526fa459cSmrg  0,                                     /* tp_as_number */
36626fa459cSmrg  0,                                     /* tp_as_sequence */
36726fa459cSmrg  0,                                     /* tp_as_mapping */
36826fa459cSmrg  0,                                     /* tp_hash  */
36926fa459cSmrg  0,                                     /* tp_call */
37026fa459cSmrg  0,                                     /* tp_str */
37126fa459cSmrg  0,                                     /* tp_getattro */
37226fa459cSmrg  0,                                     /* tp_setattro */
37326fa459cSmrg  0,                                     /* tp_as_buffer */
37426fa459cSmrg  Py_TPFLAGS_DEFAULT,                    /* tp_flags */
37526fa459cSmrg  brotli_Compressor_doc,                 /* tp_doc */
37626fa459cSmrg  0,                                     /* tp_traverse */
37726fa459cSmrg  0,                                     /* tp_clear */
37826fa459cSmrg  0,                                     /* tp_richcompare */
37926fa459cSmrg  0,                                     /* tp_weaklistoffset */
38026fa459cSmrg  0,                                     /* tp_iter */
38126fa459cSmrg  0,                                     /* tp_iternext */
38226fa459cSmrg  brotli_Compressor_methods,             /* tp_methods */
38326fa459cSmrg  brotli_Compressor_members,             /* tp_members */
38426fa459cSmrg  0,                                     /* tp_getset */
38526fa459cSmrg  0,                                     /* tp_base */
38626fa459cSmrg  0,                                     /* tp_dict */
38726fa459cSmrg  0,                                     /* tp_descr_get */
38826fa459cSmrg  0,                                     /* tp_descr_set */
38926fa459cSmrg  0,                                     /* tp_dictoffset */
39026fa459cSmrg  (initproc)brotli_Compressor_init,      /* tp_init */
39126fa459cSmrg  0,                                     /* tp_alloc */
39226fa459cSmrg  brotli_Compressor_new,                 /* tp_new */
39326fa459cSmrg};
39426fa459cSmrg
39526fa459cSmrgstatic BROTLI_BOOL decompress_stream(BrotliDecoderState* dec,
39626fa459cSmrg                                     std::vector<uint8_t>* output,
39726fa459cSmrg                                     uint8_t* input, size_t input_length) {
39826fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
39926fa459cSmrg  Py_BEGIN_ALLOW_THREADS
40026fa459cSmrg
40126fa459cSmrg  size_t available_in = input_length;
40226fa459cSmrg  const uint8_t* next_in = input;
40326fa459cSmrg  size_t available_out = 0;
40426fa459cSmrg  uint8_t* next_out = NULL;
40526fa459cSmrg
40626fa459cSmrg  BrotliDecoderResult result = BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT;
40726fa459cSmrg  while (result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) {
40826fa459cSmrg    result = BrotliDecoderDecompressStream(dec,
40926fa459cSmrg                                           &available_in, &next_in,
41026fa459cSmrg                                           &available_out, &next_out, NULL);
41126fa459cSmrg    size_t buffer_length = 0; // Request all available output.
41226fa459cSmrg    const uint8_t* buffer = BrotliDecoderTakeOutput(dec, &buffer_length);
41326fa459cSmrg    if (buffer_length) {
41426fa459cSmrg      (*output).insert((*output).end(), buffer, buffer + buffer_length);
41526fa459cSmrg    }
41626fa459cSmrg  }
41726fa459cSmrg  ok = result != BROTLI_DECODER_RESULT_ERROR && !available_in;
41826fa459cSmrg
41926fa459cSmrg  Py_END_ALLOW_THREADS
42026fa459cSmrg  return ok;
42126fa459cSmrg}
42226fa459cSmrg
42326fa459cSmrgPyDoc_STRVAR(brotli_Decompressor_doc,
42426fa459cSmrg"An object to decompress a byte string.\n"
42526fa459cSmrg"\n"
42626fa459cSmrg"Signature:\n"
42726fa459cSmrg"  Decompressor()\n"
42826fa459cSmrg"\n"
42926fa459cSmrg"Raises:\n"
43026fa459cSmrg"  brotli.error: If arguments are invalid.\n");
43126fa459cSmrg
43226fa459cSmrgtypedef struct {
43326fa459cSmrg  PyObject_HEAD
43426fa459cSmrg  BrotliDecoderState* dec;
43526fa459cSmrg} brotli_Decompressor;
43626fa459cSmrg
43726fa459cSmrgstatic void brotli_Decompressor_dealloc(brotli_Decompressor* self) {
43826fa459cSmrg  BrotliDecoderDestroyInstance(self->dec);
43926fa459cSmrg  #if PY_MAJOR_VERSION >= 3
44026fa459cSmrg  Py_TYPE(self)->tp_free((PyObject*)self);
44126fa459cSmrg  #else
44226fa459cSmrg  self->ob_type->tp_free((PyObject*)self);
44326fa459cSmrg  #endif
44426fa459cSmrg}
44526fa459cSmrg
44626fa459cSmrgstatic PyObject* brotli_Decompressor_new(PyTypeObject *type, PyObject *args, PyObject *keywds) {
44726fa459cSmrg  brotli_Decompressor *self;
44826fa459cSmrg  self = (brotli_Decompressor *)type->tp_alloc(type, 0);
44926fa459cSmrg
45026fa459cSmrg  if (self != NULL) {
45126fa459cSmrg    self->dec = BrotliDecoderCreateInstance(0, 0, 0);
45226fa459cSmrg  }
45326fa459cSmrg
45426fa459cSmrg  return (PyObject *)self;
45526fa459cSmrg}
45626fa459cSmrg
45726fa459cSmrgstatic int brotli_Decompressor_init(brotli_Decompressor *self, PyObject *args, PyObject *keywds) {
45826fa459cSmrg  int ok;
45926fa459cSmrg
46026fa459cSmrg  static const char *kwlist[] = {NULL};
46126fa459cSmrg
46226fa459cSmrg  ok = PyArg_ParseTupleAndKeywords(args, keywds, "|:Decompressor",
46326fa459cSmrg                    const_cast<char **>(kwlist));
46426fa459cSmrg  if (!ok)
46526fa459cSmrg    return -1;
46626fa459cSmrg  if (!self->dec)
46726fa459cSmrg    return -1;
46826fa459cSmrg
46926fa459cSmrg  return 0;
47026fa459cSmrg}
47126fa459cSmrg
47226fa459cSmrgPyDoc_STRVAR(brotli_Decompressor_process_doc,
47326fa459cSmrg"Process \"string\" for decompression, returning a string that contains \n"
47426fa459cSmrg"decompressed output data.  This data should be concatenated to the output \n"
47526fa459cSmrg"produced by any preceding calls to the \"process()\" method. \n"
47626fa459cSmrg"Some or all of the input may be kept in internal buffers for later \n"
47726fa459cSmrg"processing, and the decompressed output data may be empty until enough input \n"
47826fa459cSmrg"has been accumulated.\n"
47926fa459cSmrg"\n"
48026fa459cSmrg"Signature:\n"
48126fa459cSmrg"  decompress(string)\n"
48226fa459cSmrg"\n"
48326fa459cSmrg"Args:\n"
48426fa459cSmrg"  string (bytes): The input data\n"
48526fa459cSmrg"\n"
48626fa459cSmrg"Returns:\n"
48726fa459cSmrg"  The decompressed output data (bytes)\n"
48826fa459cSmrg"\n"
48926fa459cSmrg"Raises:\n"
49026fa459cSmrg"  brotli.error: If decompression fails\n");
49126fa459cSmrg
49226fa459cSmrgstatic PyObject* brotli_Decompressor_process(brotli_Decompressor *self, PyObject *args) {
49326fa459cSmrg  PyObject* ret = NULL;
49426fa459cSmrg  std::vector<uint8_t> output;
49526fa459cSmrg  Py_buffer input;
49626fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
49726fa459cSmrg
49826fa459cSmrg#if PY_MAJOR_VERSION >= 3
49926fa459cSmrg  ok = (BROTLI_BOOL)PyArg_ParseTuple(args, "y*:process", &input);
50026fa459cSmrg#else
50126fa459cSmrg  ok = (BROTLI_BOOL)PyArg_ParseTuple(args, "s*:process", &input);
50226fa459cSmrg#endif
50326fa459cSmrg
50426fa459cSmrg  if (!ok)
50526fa459cSmrg    return NULL;
50626fa459cSmrg
50726fa459cSmrg  if (!self->dec) {
50826fa459cSmrg    ok = BROTLI_FALSE;
50926fa459cSmrg    goto end;
51026fa459cSmrg  }
51126fa459cSmrg
51226fa459cSmrg  ok = decompress_stream(self->dec, &output, static_cast<uint8_t*>(input.buf), input.len);
51326fa459cSmrg
51426fa459cSmrgend:
51526fa459cSmrg  PyBuffer_Release(&input);
51626fa459cSmrg  if (ok) {
51726fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.empty() ? NULL : &output[0]), output.size());
51826fa459cSmrg  } else {
51926fa459cSmrg    PyErr_SetString(BrotliError, "BrotliDecoderDecompressStream failed while processing the stream");
52026fa459cSmrg  }
52126fa459cSmrg
52226fa459cSmrg  return ret;
52326fa459cSmrg}
52426fa459cSmrg
52526fa459cSmrgPyDoc_STRVAR(brotli_Decompressor_is_finished_doc,
52626fa459cSmrg"Checks if decoder instance reached the final state.\n"
52726fa459cSmrg"\n"
52826fa459cSmrg"Signature:\n"
52926fa459cSmrg"  is_finished()\n"
53026fa459cSmrg"\n"
53126fa459cSmrg"Returns:\n"
53226fa459cSmrg"  True  if the decoder is in a state where it reached the end of the input\n"
53326fa459cSmrg"        and produced all of the output\n"
53426fa459cSmrg"  False otherwise\n"
53526fa459cSmrg"\n"
53626fa459cSmrg"Raises:\n"
53726fa459cSmrg"  brotli.error: If decompression fails\n");
53826fa459cSmrg
53926fa459cSmrgstatic PyObject* brotli_Decompressor_is_finished(brotli_Decompressor *self) {
54026fa459cSmrg  PyObject *ret = NULL;
54126fa459cSmrg  std::vector<uint8_t> output;
54226fa459cSmrg  BROTLI_BOOL ok = BROTLI_TRUE;
54326fa459cSmrg
54426fa459cSmrg  if (!self->dec) {
54526fa459cSmrg    ok = BROTLI_FALSE;
54626fa459cSmrg    PyErr_SetString(BrotliError, "BrotliDecoderState is NULL while checking is_finished");
54726fa459cSmrg    goto end;
54826fa459cSmrg  }
54926fa459cSmrg
55026fa459cSmrg  if (BrotliDecoderIsFinished(self->dec)) {
55126fa459cSmrg    Py_RETURN_TRUE;
55226fa459cSmrg  } else {
55326fa459cSmrg    Py_RETURN_FALSE;
55426fa459cSmrg  }
55526fa459cSmrg
55626fa459cSmrgend:
55726fa459cSmrg  if (ok) {
55826fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.empty() ? NULL : &output[0]), output.size());
55926fa459cSmrg  } else {
56026fa459cSmrg    PyErr_SetString(BrotliError, "BrotliDecoderDecompressStream failed while finishing the stream");
56126fa459cSmrg  }
56226fa459cSmrg
56326fa459cSmrg  return ret;
56426fa459cSmrg}
56526fa459cSmrg
56626fa459cSmrgstatic PyMemberDef brotli_Decompressor_members[] = {
56726fa459cSmrg  {NULL}  /* Sentinel */
56826fa459cSmrg};
56926fa459cSmrg
57026fa459cSmrgstatic PyMethodDef brotli_Decompressor_methods[] = {
57126fa459cSmrg  {"process", (PyCFunction)brotli_Decompressor_process, METH_VARARGS, brotli_Decompressor_process_doc},
57226fa459cSmrg  {"is_finished", (PyCFunction)brotli_Decompressor_is_finished, METH_NOARGS, brotli_Decompressor_is_finished_doc},
57326fa459cSmrg  {NULL}  /* Sentinel */
57426fa459cSmrg};
57526fa459cSmrg
57626fa459cSmrgstatic PyTypeObject brotli_DecompressorType = {
57726fa459cSmrg  #if PY_MAJOR_VERSION >= 3
57826fa459cSmrg  PyVarObject_HEAD_INIT(NULL, 0)
57926fa459cSmrg  #else
58026fa459cSmrg  PyObject_HEAD_INIT(NULL)
58126fa459cSmrg  0,                                     /* ob_size*/
58226fa459cSmrg  #endif
58326fa459cSmrg  "brotli.Decompressor",                   /* tp_name */
58426fa459cSmrg  sizeof(brotli_Decompressor),             /* tp_basicsize */
58526fa459cSmrg  0,                                       /* tp_itemsize */
58626fa459cSmrg  (destructor)brotli_Decompressor_dealloc, /* tp_dealloc */
58726fa459cSmrg  0,                                       /* tp_print */
58826fa459cSmrg  0,                                       /* tp_getattr */
58926fa459cSmrg  0,                                       /* tp_setattr */
59026fa459cSmrg  0,                                       /* tp_compare */
59126fa459cSmrg  0,                                       /* tp_repr */
59226fa459cSmrg  0,                                       /* tp_as_number */
59326fa459cSmrg  0,                                       /* tp_as_sequence */
59426fa459cSmrg  0,                                       /* tp_as_mapping */
59526fa459cSmrg  0,                                       /* tp_hash  */
59626fa459cSmrg  0,                                       /* tp_call */
59726fa459cSmrg  0,                                       /* tp_str */
59826fa459cSmrg  0,                                       /* tp_getattro */
59926fa459cSmrg  0,                                       /* tp_setattro */
60026fa459cSmrg  0,                                       /* tp_as_buffer */
60126fa459cSmrg  Py_TPFLAGS_DEFAULT,                      /* tp_flags */
60226fa459cSmrg  brotli_Decompressor_doc,                 /* tp_doc */
60326fa459cSmrg  0,                                       /* tp_traverse */
60426fa459cSmrg  0,                                       /* tp_clear */
60526fa459cSmrg  0,                                       /* tp_richcompare */
60626fa459cSmrg  0,                                       /* tp_weaklistoffset */
60726fa459cSmrg  0,                                       /* tp_iter */
60826fa459cSmrg  0,                                       /* tp_iternext */
60926fa459cSmrg  brotli_Decompressor_methods,             /* tp_methods */
61026fa459cSmrg  brotli_Decompressor_members,             /* tp_members */
61126fa459cSmrg  0,                                       /* tp_getset */
61226fa459cSmrg  0,                                       /* tp_base */
61326fa459cSmrg  0,                                       /* tp_dict */
61426fa459cSmrg  0,                                       /* tp_descr_get */
61526fa459cSmrg  0,                                       /* tp_descr_set */
61626fa459cSmrg  0,                                       /* tp_dictoffset */
61726fa459cSmrg  (initproc)brotli_Decompressor_init,      /* tp_init */
61826fa459cSmrg  0,                                       /* tp_alloc */
61926fa459cSmrg  brotli_Decompressor_new,                 /* tp_new */
62026fa459cSmrg};
62126fa459cSmrg
62226fa459cSmrgPyDoc_STRVAR(brotli_decompress__doc__,
62326fa459cSmrg"Decompress a compressed byte string.\n"
62426fa459cSmrg"\n"
62526fa459cSmrg"Signature:\n"
62626fa459cSmrg"  decompress(string)\n"
62726fa459cSmrg"\n"
62826fa459cSmrg"Args:\n"
62926fa459cSmrg"  string (bytes): The compressed input data.\n"
63026fa459cSmrg"\n"
63126fa459cSmrg"Returns:\n"
63226fa459cSmrg"  The decompressed byte string.\n"
63326fa459cSmrg"\n"
63426fa459cSmrg"Raises:\n"
63526fa459cSmrg"  brotli.error: If decompressor fails.\n");
63626fa459cSmrg
63726fa459cSmrgstatic PyObject* brotli_decompress(PyObject *self, PyObject *args, PyObject *keywds) {
63826fa459cSmrg  PyObject *ret = NULL;
63926fa459cSmrg  Py_buffer input;
64026fa459cSmrg  const uint8_t* next_in;
64126fa459cSmrg  size_t available_in;
64226fa459cSmrg  int ok;
64326fa459cSmrg
64426fa459cSmrg  static const char *kwlist[] = {"string", NULL};
64526fa459cSmrg
64626fa459cSmrg#if PY_MAJOR_VERSION >= 3
64726fa459cSmrg  ok = PyArg_ParseTupleAndKeywords(args, keywds, "y*|:decompress",
64826fa459cSmrg                                   const_cast<char **>(kwlist), &input);
64926fa459cSmrg#else
65026fa459cSmrg  ok = PyArg_ParseTupleAndKeywords(args, keywds, "s*|:decompress",
65126fa459cSmrg                                   const_cast<char **>(kwlist), &input);
65226fa459cSmrg#endif
65326fa459cSmrg
65426fa459cSmrg  if (!ok)
65526fa459cSmrg    return NULL;
65626fa459cSmrg
65726fa459cSmrg  std::vector<uint8_t> output;
65826fa459cSmrg
65926fa459cSmrg  /* >>> Pure C block; release python GIL. */
66026fa459cSmrg  Py_BEGIN_ALLOW_THREADS
66126fa459cSmrg
66226fa459cSmrg  BrotliDecoderState* state = BrotliDecoderCreateInstance(0, 0, 0);
66326fa459cSmrg
66426fa459cSmrg  BrotliDecoderResult result = BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT;
66526fa459cSmrg  next_in = static_cast<uint8_t*>(input.buf);
66626fa459cSmrg  available_in = input.len;
66726fa459cSmrg  while (result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) {
66826fa459cSmrg    size_t available_out = 0;
66926fa459cSmrg    result = BrotliDecoderDecompressStream(state, &available_in, &next_in,
67026fa459cSmrg                                           &available_out, 0, 0);
67126fa459cSmrg    const uint8_t* next_out = BrotliDecoderTakeOutput(state, &available_out);
67226fa459cSmrg    if (available_out != 0)
67326fa459cSmrg      output.insert(output.end(), next_out, next_out + available_out);
67426fa459cSmrg  }
67526fa459cSmrg  ok = result == BROTLI_DECODER_RESULT_SUCCESS && !available_in;
67626fa459cSmrg  BrotliDecoderDestroyInstance(state);
67726fa459cSmrg
67826fa459cSmrg  Py_END_ALLOW_THREADS
67926fa459cSmrg  /* <<< Pure C block end. Python GIL reacquired. */
68026fa459cSmrg
68126fa459cSmrg  PyBuffer_Release(&input);
68226fa459cSmrg  if (ok) {
68326fa459cSmrg    ret = PyBytes_FromStringAndSize((char*)(output.size() ? &output[0] : NULL), output.size());
68426fa459cSmrg  } else {
68526fa459cSmrg    PyErr_SetString(BrotliError, "BrotliDecompress failed");
68626fa459cSmrg  }
68726fa459cSmrg
68826fa459cSmrg  return ret;
68926fa459cSmrg}
69026fa459cSmrg
69126fa459cSmrgstatic PyMethodDef brotli_methods[] = {
69226fa459cSmrg  {"decompress", (PyCFunction)brotli_decompress, METH_VARARGS | METH_KEYWORDS, brotli_decompress__doc__},
69326fa459cSmrg  {NULL, NULL, 0, NULL}
69426fa459cSmrg};
69526fa459cSmrg
69626fa459cSmrgPyDoc_STRVAR(brotli_doc, "Implementation module for the Brotli library.");
69726fa459cSmrg
69826fa459cSmrg#if PY_MAJOR_VERSION >= 3
69926fa459cSmrg#define INIT_BROTLI   PyInit__brotli
70026fa459cSmrg#define CREATE_BROTLI PyModule_Create(&brotli_module)
70126fa459cSmrg#define RETURN_BROTLI return m
70226fa459cSmrg#define RETURN_NULL return NULL
70326fa459cSmrg
70426fa459cSmrgstatic struct PyModuleDef brotli_module = {
70526fa459cSmrg  PyModuleDef_HEAD_INIT,
70626fa459cSmrg  "_brotli",      /* m_name */
70726fa459cSmrg  brotli_doc,     /* m_doc */
70826fa459cSmrg  0,              /* m_size */
70926fa459cSmrg  brotli_methods, /* m_methods */
71026fa459cSmrg  NULL,           /* m_reload */
71126fa459cSmrg  NULL,           /* m_traverse */
71226fa459cSmrg  NULL,           /* m_clear */
71326fa459cSmrg  NULL            /* m_free */
71426fa459cSmrg};
71526fa459cSmrg#else
71626fa459cSmrg#define INIT_BROTLI   init_brotli
71726fa459cSmrg#define CREATE_BROTLI Py_InitModule3("_brotli", brotli_methods, brotli_doc)
71826fa459cSmrg#define RETURN_BROTLI return
71926fa459cSmrg#define RETURN_NULL return
72026fa459cSmrg#endif
72126fa459cSmrg
72226fa459cSmrgPyMODINIT_FUNC INIT_BROTLI(void) {
72326fa459cSmrg  PyObject *m = CREATE_BROTLI;
72426fa459cSmrg
72526fa459cSmrg  BrotliError = PyErr_NewException((char*) "brotli.error", NULL, NULL);
72626fa459cSmrg  if (BrotliError != NULL) {
72726fa459cSmrg    Py_INCREF(BrotliError);
72826fa459cSmrg    PyModule_AddObject(m, "error", BrotliError);
72926fa459cSmrg  }
73026fa459cSmrg
73126fa459cSmrg  if (PyType_Ready(&brotli_CompressorType) < 0) {
73226fa459cSmrg    RETURN_NULL;
73326fa459cSmrg  }
73426fa459cSmrg  Py_INCREF(&brotli_CompressorType);
73526fa459cSmrg  PyModule_AddObject(m, "Compressor", (PyObject *)&brotli_CompressorType);
73626fa459cSmrg
73726fa459cSmrg  if (PyType_Ready(&brotli_DecompressorType) < 0) {
73826fa459cSmrg    RETURN_NULL;
73926fa459cSmrg  }
74026fa459cSmrg  Py_INCREF(&brotli_DecompressorType);
74126fa459cSmrg  PyModule_AddObject(m, "Decompressor", (PyObject *)&brotli_DecompressorType);
74226fa459cSmrg
74326fa459cSmrg  PyModule_AddIntConstant(m, "MODE_GENERIC", (int) BROTLI_MODE_GENERIC);
74426fa459cSmrg  PyModule_AddIntConstant(m, "MODE_TEXT", (int) BROTLI_MODE_TEXT);
74526fa459cSmrg  PyModule_AddIntConstant(m, "MODE_FONT", (int) BROTLI_MODE_FONT);
74626fa459cSmrg
74726fa459cSmrg  char version[16];
74826fa459cSmrg  snprintf(version, sizeof(version), "%d.%d.%d",
74926fa459cSmrg      BROTLI_VERSION >> 24, (BROTLI_VERSION >> 12) & 0xFFF, BROTLI_VERSION & 0xFFF);
75026fa459cSmrg  PyModule_AddStringConstant(m, "__version__", version);
75126fa459cSmrg
75226fa459cSmrg  RETURN_BROTLI;
75326fa459cSmrg}
754