Mercurial > hg
view contrib/python-zstandard/c-ext/compressiondict.c @ 50329:3dbc7b1ecaba stable
typing: correct the signature of error.CommandError
There's a place in `mercurial.dispatch._parse()` that passes None if a parse
error happens before the command can be parsed out, and casting the error to
bytes works fine because the command and message fields are apparently ignored.
Likewise, TortoiseHg similarly passes None for the same reason.
author | Matt Harbison <matt_harbison@yahoo.com> |
---|---|
date | Fri, 24 Mar 2023 02:22:12 -0400 |
parents | e92ca942ddca |
children |
line wrap: on
line source
/** * Copyright (c) 2016-present, Gregory Szorc * All rights reserved. * * This software may be modified and distributed under the terms * of the BSD license. See the LICENSE file for details. */ #include "python-zstandard.h" extern PyObject* ZstdError; ZstdCompressionDict* train_dictionary(PyObject* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "dict_size", "samples", "k", "d", "notifications", "dict_id", "level", "steps", "threads", NULL }; size_t capacity; PyObject* samples; unsigned k = 0; unsigned d = 0; unsigned notifications = 0; unsigned dictID = 0; int level = 0; unsigned steps = 0; int threads = 0; ZDICT_cover_params_t params; Py_ssize_t samplesLen; Py_ssize_t i; size_t samplesSize = 0; void* sampleBuffer = NULL; size_t* sampleSizes = NULL; void* sampleOffset; Py_ssize_t sampleSize; void* dict = NULL; size_t zresult; ZstdCompressionDict* result = NULL; if (!PyArg_ParseTupleAndKeywords(args, kwargs, "nO!|IIIIiIi:train_dictionary", kwlist, &capacity, &PyList_Type, &samples, &k, &d, ¬ifications, &dictID, &level, &steps, &threads)) { return NULL; } if (threads < 0) { threads = cpu_count(); } memset(¶ms, 0, sizeof(params)); params.k = k; params.d = d; params.steps = steps; params.nbThreads = threads; params.zParams.notificationLevel = notifications; params.zParams.dictID = dictID; params.zParams.compressionLevel = level; /* Figure out total size of input samples. */ samplesLen = PyList_Size(samples); for (i = 0; i < samplesLen; i++) { PyObject* sampleItem = PyList_GET_ITEM(samples, i); if (!PyBytes_Check(sampleItem)) { PyErr_SetString(PyExc_ValueError, "samples must be bytes"); return NULL; } samplesSize += PyBytes_GET_SIZE(sampleItem); } sampleBuffer = PyMem_Malloc(samplesSize); if (!sampleBuffer) { PyErr_NoMemory(); goto finally; } sampleSizes = PyMem_Malloc(samplesLen * sizeof(size_t)); if (!sampleSizes) { PyErr_NoMemory(); goto finally; } sampleOffset = sampleBuffer; for (i = 0; i < samplesLen; i++) { PyObject* sampleItem = PyList_GET_ITEM(samples, i); sampleSize = PyBytes_GET_SIZE(sampleItem); sampleSizes[i] = sampleSize; memcpy(sampleOffset, PyBytes_AS_STRING(sampleItem), sampleSize); sampleOffset = (char*)sampleOffset + sampleSize; } dict = PyMem_Malloc(capacity); if (!dict) { PyErr_NoMemory(); goto finally; } Py_BEGIN_ALLOW_THREADS /* No parameters uses the default function, which will use default params and call ZDICT_optimizeTrainFromBuffer_cover under the hood. */ if (!params.k && !params.d && !params.zParams.compressionLevel && !params.zParams.notificationLevel && !params.zParams.dictID) { zresult = ZDICT_trainFromBuffer(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen); } /* Use optimize mode if user controlled steps or threads explicitly. */ else if (params.steps || params.nbThreads) { zresult = ZDICT_optimizeTrainFromBuffer_cover(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen, ¶ms); } /* Non-optimize mode with explicit control. */ else { zresult = ZDICT_trainFromBuffer_cover(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen, params); } Py_END_ALLOW_THREADS if (ZDICT_isError(zresult)) { PyMem_Free(dict); PyErr_Format(ZstdError, "cannot train dict: %s", ZDICT_getErrorName(zresult)); goto finally; } result = PyObject_New(ZstdCompressionDict, &ZstdCompressionDictType); if (!result) { PyMem_Free(dict); goto finally; } result->dictData = dict; result->dictSize = zresult; result->dictType = ZSTD_dct_fullDict; result->d = params.d; result->k = params.k; result->cdict = NULL; result->ddict = NULL; finally: PyMem_Free(sampleBuffer); PyMem_Free(sampleSizes); return result; } int ensure_ddict(ZstdCompressionDict* dict) { if (dict->ddict) { return 0; } Py_BEGIN_ALLOW_THREADS dict->ddict = ZSTD_createDDict_advanced(dict->dictData, dict->dictSize, ZSTD_dlm_byRef, dict->dictType, ZSTD_defaultCMem); Py_END_ALLOW_THREADS if (!dict->ddict) { PyErr_SetString(ZstdError, "could not create decompression dict"); return 1; } return 0; } PyDoc_STRVAR(ZstdCompressionDict__doc__, "ZstdCompressionDict(data) - Represents a computed compression dictionary\n" "\n" "This type holds the results of a computed Zstandard compression dictionary.\n" "Instances are obtained by calling ``train_dictionary()`` or by passing\n" "bytes obtained from another source into the constructor.\n" ); static int ZstdCompressionDict_init(ZstdCompressionDict* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "data", "dict_type", NULL }; int result = -1; Py_buffer source; unsigned dictType = ZSTD_dct_auto; self->dictData = NULL; self->dictSize = 0; self->cdict = NULL; self->ddict = NULL; #if PY_MAJOR_VERSION >= 3 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|I:ZstdCompressionDict", #else if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*|I:ZstdCompressionDict", #endif kwlist, &source, &dictType)) { return -1; } if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) { PyErr_SetString(PyExc_ValueError, "data buffer should be contiguous and have at most one dimension"); goto finally; } if (dictType != ZSTD_dct_auto && dictType != ZSTD_dct_rawContent && dictType != ZSTD_dct_fullDict) { PyErr_Format(PyExc_ValueError, "invalid dictionary load mode: %d; must use DICT_TYPE_* constants", dictType); goto finally; } self->dictType = dictType; self->dictData = PyMem_Malloc(source.len); if (!self->dictData) { PyErr_NoMemory(); goto finally; } memcpy(self->dictData, source.buf, source.len); self->dictSize = source.len; result = 0; finally: PyBuffer_Release(&source); return result; } static void ZstdCompressionDict_dealloc(ZstdCompressionDict* self) { if (self->cdict) { ZSTD_freeCDict(self->cdict); self->cdict = NULL; } if (self->ddict) { ZSTD_freeDDict(self->ddict); self->ddict = NULL; } if (self->dictData) { PyMem_Free(self->dictData); self->dictData = NULL; } PyObject_Del(self); } PyDoc_STRVAR(ZstdCompressionDict_precompute_compress__doc__, "Precompute a dictionary so it can be used by multiple compressors.\n" ); static PyObject* ZstdCompressionDict_precompute_compress(ZstdCompressionDict* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "level", "compression_params", NULL }; int level = 0; ZstdCompressionParametersObject* compressionParams = NULL; ZSTD_compressionParameters cParams; size_t zresult; if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!:precompute_compress", kwlist, &level, &ZstdCompressionParametersType, &compressionParams)) { return NULL; } if (level && compressionParams) { PyErr_SetString(PyExc_ValueError, "must only specify one of level or compression_params"); return NULL; } if (!level && !compressionParams) { PyErr_SetString(PyExc_ValueError, "must specify one of level or compression_params"); return NULL; } if (self->cdict) { zresult = ZSTD_freeCDict(self->cdict); self->cdict = NULL; if (ZSTD_isError(zresult)) { PyErr_Format(ZstdError, "unable to free CDict: %s", ZSTD_getErrorName(zresult)); return NULL; } } if (level) { cParams = ZSTD_getCParams(level, 0, self->dictSize); } else { if (to_cparams(compressionParams, &cParams)) { return NULL; } } assert(!self->cdict); self->cdict = ZSTD_createCDict_advanced(self->dictData, self->dictSize, ZSTD_dlm_byRef, self->dictType, cParams, ZSTD_defaultCMem); if (!self->cdict) { PyErr_SetString(ZstdError, "unable to precompute dictionary"); return NULL; } Py_RETURN_NONE; } static PyObject* ZstdCompressionDict_dict_id(ZstdCompressionDict* self) { unsigned dictID = ZDICT_getDictID(self->dictData, self->dictSize); return PyLong_FromLong(dictID); } static PyObject* ZstdCompressionDict_as_bytes(ZstdCompressionDict* self) { return PyBytes_FromStringAndSize(self->dictData, self->dictSize); } static PyMethodDef ZstdCompressionDict_methods[] = { { "dict_id", (PyCFunction)ZstdCompressionDict_dict_id, METH_NOARGS, PyDoc_STR("dict_id() -- obtain the numeric dictionary ID") }, { "as_bytes", (PyCFunction)ZstdCompressionDict_as_bytes, METH_NOARGS, PyDoc_STR("as_bytes() -- obtain the raw bytes constituting the dictionary data") }, { "precompute_compress", (PyCFunction)ZstdCompressionDict_precompute_compress, METH_VARARGS | METH_KEYWORDS, ZstdCompressionDict_precompute_compress__doc__ }, { NULL, NULL } }; static PyMemberDef ZstdCompressionDict_members[] = { { "k", T_UINT, offsetof(ZstdCompressionDict, k), READONLY, "segment size" }, { "d", T_UINT, offsetof(ZstdCompressionDict, d), READONLY, "dmer size" }, { NULL } }; static Py_ssize_t ZstdCompressionDict_length(ZstdCompressionDict* self) { return self->dictSize; } static PySequenceMethods ZstdCompressionDict_sq = { (lenfunc)ZstdCompressionDict_length, /* sq_length */ 0, /* sq_concat */ 0, /* sq_repeat */ 0, /* sq_item */ 0, /* sq_ass_item */ 0, /* sq_contains */ 0, /* sq_inplace_concat */ 0 /* sq_inplace_repeat */ }; PyTypeObject ZstdCompressionDictType = { PyVarObject_HEAD_INIT(NULL, 0) "zstd.ZstdCompressionDict", /* tp_name */ sizeof(ZstdCompressionDict), /* tp_basicsize */ 0, /* tp_itemsize */ (destructor)ZstdCompressionDict_dealloc, /* tp_dealloc */ 0, /* tp_print */ 0, /* tp_getattr */ 0, /* tp_setattr */ 0, /* tp_compare */ 0, /* tp_repr */ 0, /* tp_as_number */ &ZstdCompressionDict_sq, /* tp_as_sequence */ 0, /* tp_as_mapping */ 0, /* tp_hash */ 0, /* tp_call */ 0, /* tp_str */ 0, /* tp_getattro */ 0, /* tp_setattro */ 0, /* tp_as_buffer */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */ ZstdCompressionDict__doc__, /* tp_doc */ 0, /* tp_traverse */ 0, /* tp_clear */ 0, /* tp_richcompare */ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ ZstdCompressionDict_methods, /* tp_methods */ ZstdCompressionDict_members, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ 0, /* tp_dict */ 0, /* tp_descr_get */ 0, /* tp_descr_set */ 0, /* tp_dictoffset */ (initproc)ZstdCompressionDict_init, /* tp_init */ 0, /* tp_alloc */ PyType_GenericNew, /* tp_new */ }; void compressiondict_module_init(PyObject* mod) { Py_SET_TYPE(&ZstdCompressionDictType, &PyType_Type); if (PyType_Ready(&ZstdCompressionDictType) < 0) { return; } Py_INCREF((PyObject*)&ZstdCompressionDictType); PyModule_AddObject(mod, "ZstdCompressionDict", (PyObject*)&ZstdCompressionDictType); }