procutil: make stream detection in make_line_buffered more correct and strict
In make_line_buffered(), we don’t want to wrap the stream if we know that lines
get flushed to the underlying raw stream already.
Previously, the heuristic was too optimistic. It assumed that any stream which
is not an instance of io.BufferedIOBase doesn’t need wrapping. However, there
are buffered streams that aren’t instances of io.BufferedIOBase, like
Mercurial’s own winstdout.
The new logic is different in two ways:
First, only for the check, if unwraps any combination of WriteAllWrapper and
winstdout.
Second, it skips wrapping the stream only if it is an instance of io.RawIOBase
(or already wrapped). If it is an instance of io.BufferedIOBase, it gets
wrapped. In any other case, the function raises an exception. This ensures
that, if an unknown stream is passed or we add another wrapper in the future,
we don’t wrap the stream if it’s already line buffered or not wrap the stream
if it’s not line buffered. In fact, this was already helpful during development
of this change. Without it, I possibly would have forgot that WriteAllWrapper
needs to be ignored for the check, leading to unnecessary wrapping if stdout is
unbuffered.
The alternative would have been to always wrap unknown streams. However, I
don’t think that anyone would benefit from being less strict. We can expect
streams from the standard library to be subclassing either io.RawIOBase or
io.BufferedIOBase, so running Mercurial in the standard way should not regress
by this change. Py2exe might replace sys.stdout and sys.stderr, but that
currently breaks Mercurial anyway and also these streams don’t claim to be
interactive, so this function is not called for them.
/**
* 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;
}