Mercurial > hg
view contrib/python-zstandard/zstd.c @ 46444:6c778d20c8c2
rust: replace ToString impls with Display
ToString is automatically implementing for everything that implements
Display, and Display can avoid allocating intermediate strings.
Differential Revision: https://phab.mercurial-scm.org/D9904
author | Simon Sapin <simon.sapin@octobus.net> |
---|---|
date | Thu, 28 Jan 2021 19:21:57 +0100 |
parents | de7838053207 |
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. */ /* A Python C extension for Zstandard. */ #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN #include <Windows.h> #elif defined(__APPLE__) || defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__) #include <sys/types.h> #include <sys/sysctl.h> #endif #include "python-zstandard.h" PyObject *ZstdError; PyDoc_STRVAR(estimate_decompression_context_size__doc__, "estimate_decompression_context_size()\n" "\n" "Estimate the amount of memory allocated to a decompression context.\n" ); static PyObject* estimate_decompression_context_size(PyObject* self) { return PyLong_FromSize_t(ZSTD_estimateDCtxSize()); } PyDoc_STRVAR(frame_content_size__doc__, "frame_content_size(data)\n" "\n" "Obtain the decompressed size of a frame." ); static PyObject* frame_content_size(PyObject* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "source", NULL }; Py_buffer source; PyObject* result = NULL; unsigned long long size; #if PY_MAJOR_VERSION >= 3 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_content_size", #else if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_content_size", #endif kwlist, &source)) { return NULL; } 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; } size = ZSTD_getFrameContentSize(source.buf, source.len); if (size == ZSTD_CONTENTSIZE_ERROR) { PyErr_SetString(ZstdError, "error when determining content size"); } else if (size == ZSTD_CONTENTSIZE_UNKNOWN) { result = PyLong_FromLong(-1); } else { result = PyLong_FromUnsignedLongLong(size); } finally: PyBuffer_Release(&source); return result; } PyDoc_STRVAR(frame_header_size__doc__, "frame_header_size(data)\n" "\n" "Obtain the size of a frame header.\n" ); static PyObject* frame_header_size(PyObject* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "source", NULL }; Py_buffer source; PyObject* result = NULL; size_t zresult; #if PY_MAJOR_VERSION >= 3 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_header_size", #else if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_header_size", #endif kwlist, &source)) { return NULL; } 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; } zresult = ZSTD_frameHeaderSize(source.buf, source.len); if (ZSTD_isError(zresult)) { PyErr_Format(ZstdError, "could not determine frame header size: %s", ZSTD_getErrorName(zresult)); } else { result = PyLong_FromSize_t(zresult); } finally: PyBuffer_Release(&source); return result; } PyDoc_STRVAR(get_frame_parameters__doc__, "get_frame_parameters(data)\n" "\n" "Obtains a ``FrameParameters`` instance by parsing data.\n"); PyDoc_STRVAR(train_dictionary__doc__, "train_dictionary(dict_size, samples, k=None, d=None, steps=None,\n" " threads=None,notifications=0, dict_id=0, level=0)\n" "\n" "Train a dictionary from sample data using the COVER algorithm.\n" "\n" "A compression dictionary of size ``dict_size`` will be created from the\n" "iterable of ``samples``. The raw dictionary bytes will be returned.\n" "\n" "The COVER algorithm has 2 parameters: ``k`` and ``d``. These control the\n" "*segment size* and *dmer size*. A reasonable range for ``k`` is\n" "``[16, 2048+]``. A reasonable range for ``d`` is ``[6, 16]``.\n" "``d`` must be less than or equal to ``k``.\n" "\n" "``steps`` can be specified to control the number of steps through potential\n" "values of ``k`` and ``d`` to try. ``k`` and ``d`` will only be varied if\n" "those arguments are not defined. i.e. if ``d`` is ``8``, then only ``k``\n" "will be varied in this mode.\n" "\n" "``threads`` can specify how many threads to use to test various ``k`` and\n" "``d`` values. ``-1`` will use as many threads as available CPUs. By default,\n" "a single thread is used.\n" "\n" "When ``k`` and ``d`` are not defined, default values are used and the\n" "algorithm will perform multiple iterations - or steps - to try to find\n" "ideal parameters. If both ``k`` and ``d`` are specified, then those values\n" "will be used. ``steps`` or ``threads`` triggers optimization mode to test\n" "multiple ``k`` and ``d`` variations.\n" ); static char zstd_doc[] = "Interface to zstandard"; static PyMethodDef zstd_methods[] = { { "estimate_decompression_context_size", (PyCFunction)estimate_decompression_context_size, METH_NOARGS, estimate_decompression_context_size__doc__ }, { "frame_content_size", (PyCFunction)frame_content_size, METH_VARARGS | METH_KEYWORDS, frame_content_size__doc__ }, { "frame_header_size", (PyCFunction)frame_header_size, METH_VARARGS | METH_KEYWORDS, frame_header_size__doc__ }, { "get_frame_parameters", (PyCFunction)get_frame_parameters, METH_VARARGS | METH_KEYWORDS, get_frame_parameters__doc__ }, { "train_dictionary", (PyCFunction)train_dictionary, METH_VARARGS | METH_KEYWORDS, train_dictionary__doc__ }, { NULL, NULL } }; void bufferutil_module_init(PyObject* mod); void compressobj_module_init(PyObject* mod); void compressor_module_init(PyObject* mod); void compressionparams_module_init(PyObject* mod); void constants_module_init(PyObject* mod); void compressionchunker_module_init(PyObject* mod); void compressiondict_module_init(PyObject* mod); void compressionreader_module_init(PyObject* mod); void compressionwriter_module_init(PyObject* mod); void compressoriterator_module_init(PyObject* mod); void decompressor_module_init(PyObject* mod); void decompressobj_module_init(PyObject* mod); void decompressionreader_module_init(PyObject *mod); void decompressionwriter_module_init(PyObject* mod); void decompressoriterator_module_init(PyObject* mod); void frameparams_module_init(PyObject* mod); void zstd_module_init(PyObject* m) { /* python-zstandard relies on unstable zstd C API features. This means that changes in zstd may break expectations in python-zstandard. python-zstandard is distributed with a copy of the zstd sources. python-zstandard is only guaranteed to work with the bundled version of zstd. However, downstream redistributors or packagers may unbundle zstd from python-zstandard. This can result in a mismatch between zstd versions and API semantics. This essentially "voids the warranty" of python-zstandard and may cause undefined behavior. We detect this mismatch here and refuse to load the module if this scenario is detected. */ if (ZSTD_VERSION_NUMBER != 10404 || ZSTD_versionNumber() != 10404) { PyErr_SetString(PyExc_ImportError, "zstd C API mismatch; Python bindings not compiled against expected zstd version"); return; } bufferutil_module_init(m); compressionparams_module_init(m); compressiondict_module_init(m); compressobj_module_init(m); compressor_module_init(m); compressionchunker_module_init(m); compressionreader_module_init(m); compressionwriter_module_init(m); compressoriterator_module_init(m); constants_module_init(m); decompressor_module_init(m); decompressobj_module_init(m); decompressionreader_module_init(m); decompressionwriter_module_init(m); decompressoriterator_module_init(m); frameparams_module_init(m); } #if defined(__GNUC__) && (__GNUC__ >= 4) # define PYTHON_ZSTD_VISIBILITY __attribute__ ((visibility ("default"))) #else # define PYTHON_ZSTD_VISIBILITY #endif #if PY_MAJOR_VERSION >= 3 static struct PyModuleDef zstd_module = { PyModuleDef_HEAD_INIT, "zstd", zstd_doc, -1, zstd_methods }; PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC PyInit_zstd(void) { PyObject *m = PyModule_Create(&zstd_module); if (m) { zstd_module_init(m); if (PyErr_Occurred()) { Py_DECREF(m); m = NULL; } } return m; } #else PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC initzstd(void) { PyObject *m = Py_InitModule3("zstd", zstd_methods, zstd_doc); if (m) { zstd_module_init(m); } } #endif /* Attempt to resolve the number of CPUs in the system. */ int cpu_count() { int count = 0; #if defined(_WIN32) SYSTEM_INFO si; si.dwNumberOfProcessors = 0; GetSystemInfo(&si); count = si.dwNumberOfProcessors; #elif defined(__APPLE__) int num; size_t size = sizeof(int); if (0 == sysctlbyname("hw.logicalcpu", &num, &size, NULL, 0)) { count = num; } #elif defined(__linux__) count = sysconf(_SC_NPROCESSORS_ONLN); #elif defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__) int mib[2]; size_t len = sizeof(count); mib[0] = CTL_HW; mib[1] = HW_NCPU; if (0 != sysctl(mib, 2, &count, &len, NULL, 0)) { count = 0; } #elif defined(__hpux) count = mpctl(MPC_GETNUMSPUS, NULL, NULL); #endif return count; } size_t roundpow2(size_t i) { i--; i |= i >> 1; i |= i >> 2; i |= i >> 4; i |= i >> 8; i |= i >> 16; i++; return i; } /* Safer version of _PyBytes_Resize(). * * _PyBytes_Resize() only works if the refcount is 1. In some scenarios, * we can get an object with a refcount > 1, even if it was just created * with PyBytes_FromStringAndSize()! That's because (at least) CPython * pre-allocates PyBytes instances of size 1 for every possible byte value. * * If non-0 is returned, obj may or may not be NULL. */ int safe_pybytes_resize(PyObject** obj, Py_ssize_t size) { PyObject* tmp; if ((*obj)->ob_refcnt == 1) { return _PyBytes_Resize(obj, size); } tmp = PyBytes_FromStringAndSize(NULL, size); if (!tmp) { return -1; } memcpy(PyBytes_AS_STRING(tmp), PyBytes_AS_STRING(*obj), PyBytes_GET_SIZE(*obj)); Py_DECREF(*obj); *obj = tmp; return 0; }