brotli/python/brotlimodule.cc

#define PY_SSIZE_T_CLEAN 1
#include <Python.h>
#include <bytesobject.h>
#include "../enc/encode.h"
#include "../dec/decode.h"

#if PY_MAJOR_VERSION >= 3
#define PyInt_Check PyLong_Check
#define PyInt_AsLong PyLong_AsLong
#endif

using namespace brotli;

static PyObject *BrotliError;

static int mode_convertor(PyObject *o, BrotliParams::Mode *mode) {
  if (!PyInt_Check(o)) {
    PyErr_SetString(BrotliError, "Invalid mode");
    return 0;
  }

  *mode = (BrotliParams::Mode) PyInt_AsLong(o);
  if (*mode != BrotliParams::Mode::MODE_TEXT &&
      *mode != BrotliParams::Mode::MODE_FONT) {
    PyErr_SetString(BrotliError, "Invalid mode");
    return 0;
  }

  return 1;
}

PyDoc_STRVAR(compress__doc__,
"compress(data[, mode=MODE_TEXT, enable_transforms=False])\n"
"\n"
"Args:\n"
"  data (bytes): The input data.\n"
"  mode (int, optional): The compression mode can be MODE_TEXT (default) or\n"
"    MODE_FONT.\n"
"  enable_dictionary (bool, optional): Controls whether to enable encoder\n"
"    dictionary. Defaults to True. Respected only if quality > 9.\n"
"  enable_transforms (bool, optional): Controls whether to enable encoder\n"
"    transforms. Defaults to False. Respected only if quality > 9.\n"
"  greedy_block_split (bool, optional): Controls whether to enable a faster\n"
"    but less dense compression mode. Defaults to False. Respected only if\n"
"    quality > 9.\n"
"\n"
"Returns:\n"
"  The compressed string of bytes.\n"
"\n"
"Raises:\n"
"  brotli.error: If mode is invalid, or compressor fails.\n");

static PyObject* brotli_compress(PyObject *self, PyObject *args, PyObject *keywds) {
  PyObject *ret = NULL;
  PyObject* enable_dictionary = NULL;
  PyObject* enable_transforms = NULL;
  PyObject* greedy_block_split = NULL;
  uint8_t *input, *output;
  size_t length, output_length;
  BrotliParams::Mode mode = (BrotliParams::Mode) -1;
  int ok;

  static char *kwlist[] = {"data", "mode", "enable_dictionary", "enable_transforms",
                           "greedy_block_split", NULL};

  ok = PyArg_ParseTupleAndKeywords(args, keywds, "s#|O&O!O!O!:compress", kwlist,
                        &input, &length,
                        &mode_convertor, &mode,
                        &PyBool_Type, &enable_dictionary,
                        &PyBool_Type, &enable_transforms,
                        &PyBool_Type, &greedy_block_split);

  if (!ok)
    return NULL;

  output_length = 1.2 * length + 10240;
  output = new uint8_t[output_length];

  BrotliParams params;
  if (mode != -1)
    params.mode = mode;
  if (enable_dictionary)
    params.enable_dictionary = PyObject_IsTrue(enable_dictionary);
  if (enable_transforms)
    params.enable_transforms = PyObject_IsTrue(enable_transforms);
  if (greedy_block_split)
    params.greedy_block_split = PyObject_IsTrue(greedy_block_split);

  ok = BrotliCompressBuffer(params, length, input,
                            &output_length, output);
  if (ok) {
    ret = PyBytes_FromStringAndSize((char*)output, output_length);
  } else {
    PyErr_SetString(BrotliError, "BrotliCompressBuffer failed");
  }

  delete[] output;

  return ret;
}

int output_callback(void* data, const uint8_t* buf, size_t count) {
  std::vector<uint8_t> *output = (std::vector<uint8_t> *)data;
  output->insert(output->end(), buf, buf + count);
  return (int)count;
}

PyDoc_STRVAR(decompress__doc__,
"decompress(string) -- Return decompressed string.");

static PyObject* brotli_decompress(PyObject *self, PyObject *args) {
  PyObject *ret = NULL;
  uint8_t *input;
  size_t length;
  int ok;

  ok = PyArg_ParseTuple(args, "s#:decompress", &input, &length);
  if (!ok)
    return NULL;

  BrotliMemInput memin;
  BrotliInput in = BrotliInitMemInput(input, length, &memin);

  BrotliOutput out;
  std::vector<uint8_t> output;
  out.cb_ = &output_callback;
  out.data_ = &output;

  ok = BrotliDecompress(in, out);
  if (ok) {
    ret = PyBytes_FromStringAndSize((char*)output.data(), output.size());
  } else {
    PyErr_SetString(BrotliError, "BrotliDecompress failed");
  }

  return ret;
}

static PyMethodDef brotli_methods[] = {
  {"compress",   (PyCFunction)brotli_compress, METH_VARARGS | METH_KEYWORDS, compress__doc__},
  {"decompress", brotli_decompress, METH_VARARGS, decompress__doc__},
  {NULL, NULL, 0, NULL}
};

PyDoc_STRVAR(brotli__doc__,
"The functions in this module allow compression and decompression using the\n"
"Brotli library.\n"
"\n"
"compress(string[, mode, transform]) -- Compress string.\n"
"decompress(string) -- Decompresses a compressed string.\n");

#if PY_MAJOR_VERSION >= 3
#define INIT_BROTLI   PyInit_brotli
#define CREATE_BROTLI PyModule_Create(&brotli_module)
#define RETURN_BROTLI return m

static struct PyModuleDef brotli_module = {
  PyModuleDef_HEAD_INIT,
  "brotli",
  brotli__doc__,
  0,
  brotli_methods,
  NULL,
  NULL,
  NULL
};
#else
#define INIT_BROTLI   initbrotli
#define CREATE_BROTLI Py_InitModule3("brotli", brotli_methods, brotli__doc__)
#define RETURN_BROTLI return
#endif

PyMODINIT_FUNC INIT_BROTLI(void) {
  PyObject *m = CREATE_BROTLI;

  BrotliError = PyErr_NewException((char*) "brotli.error", NULL, NULL);

  if (BrotliError != NULL) {
    Py_INCREF(BrotliError);
    PyModule_AddObject(m, "error", BrotliError);
  }

  PyModule_AddIntConstant(m, "MODE_TEXT", (int) BrotliParams::Mode::MODE_TEXT);
  PyModule_AddIntConstant(m, "MODE_FONT", (int) BrotliParams::Mode::MODE_FONT);

  RETURN_BROTLI;
}