Skip to content

Commit 388b69f

Browse files
emmatypingemmatyping
committed
Expand module state macros to improve code quality
Also removes module state references from the classes in the _zstd module and instead uses PyType_GetModuleState()
1 parent 68f8009 commit 388b69f

File tree

5 files changed

+179
-163
lines changed

5 files changed

+179
-163
lines changed

‎Modules/_zstd/_zstdmodule.c‎

Lines changed: 108 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -256,8 +256,10 @@ _zstd__train_dict_impl(PyObject *module, PyBytesObject *samples_bytes,
256256

257257
/* Check zstd dict error */
258258
if (ZDICT_isError(zstd_ret)) {
259-
STATE_FROM_MODULE(module);
260-
set_zstd_error(MODULE_STATE, ERR_TRAIN_DICT, zstd_ret);
259+
_zstd_state* const _module_state = PyModule_GetState(module);
260+
if (_module_state != NULL) {
261+
set_zstd_error(_module_state, ERR_TRAIN_DICT, zstd_ret);
262+
}
261263
goto error;
262264
}
263265

@@ -384,8 +386,10 @@ _zstd__finalize_dict_impl(PyObject *module, PyBytesObject *custom_dict_bytes,
384386

385387
/* Check zstd dict error */
386388
if (ZDICT_isError(zstd_ret)) {
387-
STATE_FROM_MODULE(module);
388-
set_zstd_error(MODULE_STATE, ERR_FINALIZE_DICT, zstd_ret);
389+
_zstd_state* const _module_state = PyModule_GetState(module);
390+
if (_module_state != NULL) {
391+
set_zstd_error(_module_state, ERR_FINALIZE_DICT, zstd_ret);
392+
}
389393
goto error;
390394
}
391395

@@ -425,15 +429,19 @@ _zstd__get_param_bounds_impl(PyObject *module, int is_compress,
425429
if (is_compress) {
426430
bound = ZSTD_cParam_getBounds(parameter);
427431
if (ZSTD_isError(bound.error)) {
428-
STATE_FROM_MODULE(module);
429-
set_zstd_error(MODULE_STATE, ERR_GET_C_BOUNDS, bound.error);
432+
_zstd_state* const _module_state = PyModule_GetState(module);
433+
if (_module_state != NULL) {
434+
set_zstd_error(_module_state, ERR_GET_C_BOUNDS, bound.error);
435+
}
430436
return NULL;
431437
}
432438
} else {
433439
bound = ZSTD_dParam_getBounds(parameter);
434440
if (ZSTD_isError(bound.error)) {
435-
STATE_FROM_MODULE(module);
436-
set_zstd_error(MODULE_STATE, ERR_GET_D_BOUNDS, bound.error);
441+
_zstd_state* const _module_state = PyModule_GetState(module);
442+
if (_module_state != NULL) {
443+
set_zstd_error(_module_state, ERR_GET_D_BOUNDS, bound.error);
444+
}
437445
return NULL;
438446
}
439447
}
@@ -462,13 +470,15 @@ _zstd_get_frame_size_impl(PyObject *module, Py_buffer *frame_buffer)
462470

463471
frame_size = ZSTD_findFrameCompressedSize(frame_buffer->buf, frame_buffer->len);
464472
if (ZSTD_isError(frame_size)) {
465-
STATE_FROM_MODULE(module);
466-
PyErr_Format(MS_MEMBER(ZstdError),
467-
"Error when finding the compressed size of a zstd frame. "
468-
"Make sure the frame_buffer argument starts from the "
469-
"beginning of a frame, and its length not less than this "
470-
"complete frame. Zstd error message: %s.",
471-
ZSTD_getErrorName(frame_size));
473+
_zstd_state* const _module_state = PyModule_GetState(module);
474+
if (_module_state != NULL) {
475+
PyErr_Format(_module_state->ZstdError,
476+
"Error when finding the compressed size of a zstd frame. "
477+
"Make sure the frame_buffer argument starts from the "
478+
"beginning of a frame, and its length not less than this "
479+
"complete frame. Zstd error message: %s.",
480+
ZSTD_getErrorName(frame_size));
481+
}
472482
goto error;
473483
}
474484

@@ -508,12 +518,14 @@ _zstd__get_frame_info_impl(PyObject *module, Py_buffer *frame_buffer)
508518
/* #define ZSTD_CONTENTSIZE_UNKNOWN (0ULL - 1)
509519
#define ZSTD_CONTENTSIZE_ERROR (0ULL - 2) */
510520
if (decompressed_size == ZSTD_CONTENTSIZE_ERROR) {
511-
STATE_FROM_MODULE(module);
512-
PyErr_SetString(MS_MEMBER(ZstdError),
513-
"Error when getting information from the header of "
514-
"a zstd frame. Make sure the frame_buffer argument "
515-
"starts from the beginning of a frame, and its length "
516-
"not less than the frame header (6~18 bytes).");
521+
_zstd_state* const _module_state = PyModule_GetState(module);
522+
if (_module_state != NULL) {
523+
PyErr_SetString(_module_state->ZstdError,
524+
"Error when getting information from the header of "
525+
"a zstd frame. Make sure the frame_buffer argument "
526+
"starts from the beginning of a frame, and its length "
527+
"not less than the frame header (6~18 bytes).");
528+
}
517529
goto error;
518530
}
519531

@@ -553,7 +565,10 @@ _zstd__set_parameter_types_impl(PyObject *module, PyObject *c_parameter_type,
553565
PyObject *d_parameter_type)
554566
/*[clinic end generated code: output=a13d4890ccbd2873 input=3e7d0d37c3a1045a]*/
555567
{
556-
STATE_FROM_MODULE(module);
568+
_zstd_state* const _module_state = PyModule_GetState(module);
569+
if (_module_state == NULL) {
570+
return NULL;
571+
}
557572

558573
if (!PyType_Check(c_parameter_type) || !PyType_Check(d_parameter_type)) {
559574
PyErr_SetString(PyExc_ValueError,
@@ -562,13 +577,13 @@ _zstd__set_parameter_types_impl(PyObject *module, PyObject *c_parameter_type,
562577
return NULL;
563578
}
564579

565-
Py_XDECREF(MS_MEMBER(CParameter_type));
580+
Py_XDECREF(_module_state->CParameter_type);
566581
Py_INCREF(c_parameter_type);
567-
MS_MEMBER(CParameter_type) = (PyTypeObject*)c_parameter_type;
582+
_module_state->CParameter_type = (PyTypeObject*) c_parameter_type;
568583

569-
Py_XDECREF(MS_MEMBER(DParameter_type));
584+
Py_XDECREF(_module_state->DParameter_type);
570585
Py_INCREF(d_parameter_type);
571-
MS_MEMBER(DParameter_type) = (PyTypeObject*)d_parameter_type;
586+
_module_state->DParameter_type = (PyTypeObject*)d_parameter_type;
572587

573588
Py_RETURN_NONE;
574589
}
@@ -598,20 +613,21 @@ _zstd_compress_impl(PyObject *module, Py_buffer *data, PyObject *level,
598613
PyObject *options, PyObject *zstd_dict)
599614
/*[clinic end generated code: output=0cca9399ca5c95cc input=e8a7c59073af923c]*/
600615
{
601-
STATE_FROM_MODULE(module);
616+
_zstd_state* const _module_state = PyModule_GetState(module);
617+
if (_module_state == NULL) {
618+
return NULL;
619+
}
602620
PyObject *ret = NULL;
603621
ZstdCompressor self = {0};
604622

605623
/* Initialize & set ZstdCompressor */
606624
self.cctx = ZSTD_createCCtx();
607625
if (self.cctx == NULL) {
608-
PyErr_SetString(MS_MEMBER(ZstdError),
626+
PyErr_SetString(_module_state->ZstdError,
609627
"Unable to create ZSTD_CCtx instance.");
610628
goto error;
611629
}
612630

613-
self.module_state = MODULE_STATE;
614-
615631
if (level != Py_None && options != Py_None) {
616632
PyErr_SetString(PyExc_RuntimeError, "Only one of level or options should be used.");
617633
return NULL;
@@ -679,20 +695,21 @@ _zstd_decompress_impl(PyObject *module, Py_buffer *data, PyObject *zstd_dict,
679695
Py_ssize_t initial_size;
680696
ZstdDecompressor self = {0};
681697
ZSTD_inBuffer in;
682-
STATE_FROM_MODULE(module);
698+
_zstd_state* const _module_state = PyModule_GetState(module);
699+
if (_module_state == NULL) {
700+
return NULL;
701+
}
683702
PyObject *ret = NULL;
684703

685704
/* Initialize & set ZstdDecompressor */
686705
self.dctx = ZSTD_createDCtx();
687706
if (self.dctx == NULL) {
688-
PyErr_SetString(MS_MEMBER(ZstdError),
707+
PyErr_SetString(_module_state->ZstdError,
689708
"Unable to create ZSTD_DCtx instance.");
690709
goto error;
691710
}
692711
self.at_frame_edge = 1;
693712

694-
self.module_state = MODULE_STATE;
695-
696713
/* Load dictionary to decompression context */
697714
if (zstd_dict != Py_None) {
698715
if (_PyZstd_load_d_dict(&self, zstd_dict) < 0) {
@@ -736,7 +753,7 @@ _zstd_decompress_impl(PyObject *module, Py_buffer *data, PyObject *zstd_dict,
736753
char *extra_msg = (Py_SIZE(ret) == 0) ? "." :
737754
", if want to output these decompressed data, use "
738755
"the ZstdDecompressor class to decompress.";
739-
PyErr_Format(MS_MEMBER(ZstdError),
756+
PyErr_Format(_module_state->ZstdError,
740757
"Decompression failed: zstd data ends in an incomplete "
741758
"frame, maybe the input data was truncated. Decompressed "
742759
"data is %zd bytes%s",
@@ -916,8 +933,8 @@ add_vars_to_module(PyObject *module)
916933

917934
#define ADD_STR_TO_STATE_MACRO(STR) \
918935
do { \
919-
MS_MEMBER(str_##STR) = PyUnicode_FromString(#STR); \
920-
if (MS_MEMBER(str_##STR) == NULL) { \
936+
_module_state->str_##STR = PyUnicode_FromString(#STR); \
937+
if (_module_state->str_##STR == NULL) { \
921938
return -1; \
922939
} \
923940
} while(0)
@@ -959,17 +976,20 @@ add_constant_to_type(PyTypeObject *type, const char *name, const long value)
959976
}
960977

961978
static int _zstd_exec(PyObject *module) {
962-
STATE_FROM_MODULE(module);
979+
_zstd_state* const _module_state = PyModule_GetState(module);
980+
if (_module_state == NULL) {
981+
return -1;
982+
}
963983

964984
/* Reusable objects & variables */
965-
MS_MEMBER(empty_bytes) = PyBytes_FromStringAndSize(NULL, 0);
966-
if (MS_MEMBER(empty_bytes) == NULL) {
985+
_module_state->empty_bytes = PyBytes_FromStringAndSize(NULL, 0);
986+
if (_module_state->empty_bytes == NULL) {
967987
return -1;
968988
}
969989

970-
MS_MEMBER(empty_readonly_memoryview) =
971-
PyMemoryView_FromMemory((char*)MODULE_STATE, 0, PyBUF_READ);
972-
if (MS_MEMBER(empty_readonly_memoryview) == NULL) {
990+
_module_state->empty_readonly_memoryview =
991+
PyMemoryView_FromMemory((char*)_module_state, 0, PyBUF_READ);
992+
if (_module_state->empty_readonly_memoryview == NULL) {
973993
return -1;
974994
}
975995

@@ -979,59 +999,59 @@ static int _zstd_exec(PyObject *module) {
979999
ADD_STR_TO_STATE_MACRO(write);
9801000
ADD_STR_TO_STATE_MACRO(flush);
9811001

982-
MS_MEMBER(CParameter_type) = NULL;
983-
MS_MEMBER(DParameter_type) = NULL;
1002+
_module_state->CParameter_type = NULL;
1003+
_module_state->DParameter_type = NULL;
9841004

9851005
/* Add variables to module */
9861006
if (add_vars_to_module(module) < 0) {
9871007
return -1;
9881008
}
9891009

9901010
/* ZstdError */
991-
MS_MEMBER(ZstdError) = PyErr_NewExceptionWithDoc(
1011+
_module_state->ZstdError = PyErr_NewExceptionWithDoc(
9921012
"_zstd.ZstdError",
9931013
"Call to the underlying zstd library failed.",
9941014
NULL, NULL);
995-
if (MS_MEMBER(ZstdError) == NULL) {
1015+
if (_module_state->ZstdError == NULL) {
9961016
return -1;
9971017
}
9981018

999-
Py_INCREF(MS_MEMBER(ZstdError));
1000-
if (PyModule_AddObject(module, "ZstdError", MS_MEMBER(ZstdError)) < 0) {
1001-
Py_DECREF(MS_MEMBER(ZstdError));
1019+
Py_INCREF(_module_state->ZstdError);
1020+
if (PyModule_AddObject(module, "ZstdError", _module_state->ZstdError) < 0) {
1021+
Py_DECREF(_module_state->ZstdError);
10021022
return -1;
10031023
}
10041024

10051025
/* ZstdDict */
10061026
if (add_type_to_module(module,
10071027
"ZstdDict",
10081028
&zstddict_type_spec,
1009-
&MS_MEMBER(ZstdDict_type)) < 0) {
1029+
&_module_state->ZstdDict_type) < 0) {
10101030
return -1;
10111031
}
10121032

10131033
// ZstdCompressor
10141034
if (add_type_to_module(module,
10151035
"ZstdCompressor",
10161036
&zstdcompressor_type_spec,
1017-
&MS_MEMBER(ZstdCompressor_type)) < 0) {
1037+
&_module_state->ZstdCompressor_type) < 0) {
10181038
return -1;
10191039
}
10201040

10211041
// Add EndDirective enum to ZstdCompressor
1022-
if (add_constant_to_type(MS_MEMBER(ZstdCompressor_type),
1042+
if (add_constant_to_type(_module_state->ZstdCompressor_type,
10231043
"CONTINUE",
10241044
ZSTD_e_continue) < 0) {
10251045
return -1;
10261046
}
10271047

1028-
if (add_constant_to_type(MS_MEMBER(ZstdCompressor_type),
1048+
if (add_constant_to_type(_module_state->ZstdCompressor_type,
10291049
"FLUSH_BLOCK",
10301050
ZSTD_e_flush) < 0) {
10311051
return -1;
10321052
}
10331053

1034-
if (add_constant_to_type(MS_MEMBER(ZstdCompressor_type),
1054+
if (add_constant_to_type(_module_state->ZstdCompressor_type,
10351055
"FLUSH_FRAME",
10361056
ZSTD_e_end) < 0) {
10371057
return -1;
@@ -1041,7 +1061,7 @@ static int _zstd_exec(PyObject *module) {
10411061
if (add_type_to_module(module,
10421062
"ZstdDecompressor",
10431063
&ZstdDecompressor_type_spec,
1044-
&MS_MEMBER(ZstdDecompressor_type)) < 0) {
1064+
&_module_state->ZstdDecompressor_type) < 0) {
10451065
return -1;
10461066
}
10471067

@@ -1051,48 +1071,54 @@ static int _zstd_exec(PyObject *module) {
10511071
static int
10521072
_zstd_traverse(PyObject *module, visitproc visit, void *arg)
10531073
{
1054-
STATE_FROM_MODULE(module);
1074+
_zstd_state* const _module_state = PyModule_GetState(module);
1075+
if (_module_state == NULL) {
1076+
return -1;
1077+
}
10551078

1056-
Py_VISIT(MS_MEMBER(empty_bytes));
1057-
Py_VISIT(MS_MEMBER(empty_readonly_memoryview));
1058-
Py_VISIT(MS_MEMBER(str_read));
1059-
Py_VISIT(MS_MEMBER(str_readinto));
1060-
Py_VISIT(MS_MEMBER(str_write));
1061-
Py_VISIT(MS_MEMBER(str_flush));
1079+
Py_VISIT(_module_state->empty_bytes);
1080+
Py_VISIT(_module_state->empty_readonly_memoryview);
1081+
Py_VISIT(_module_state->str_read);
1082+
Py_VISIT(_module_state->str_readinto);
1083+
Py_VISIT(_module_state->str_write);
1084+
Py_VISIT(_module_state->str_flush);
10621085

1063-
Py_VISIT(MS_MEMBER(ZstdDict_type));
1064-
Py_VISIT(MS_MEMBER(ZstdCompressor_type));
1086+
Py_VISIT(_module_state->ZstdDict_type);
1087+
Py_VISIT(_module_state->ZstdCompressor_type);
10651088

1066-
Py_VISIT(MS_MEMBER(ZstdDecompressor_type));
1089+
Py_VISIT(_module_state->ZstdDecompressor_type);
10671090

1068-
Py_VISIT(MS_MEMBER(ZstdError));
1091+
Py_VISIT(_module_state->ZstdError);
10691092

1070-
Py_VISIT(MS_MEMBER(CParameter_type));
1071-
Py_VISIT(MS_MEMBER(DParameter_type));
1093+
Py_VISIT(_module_state->CParameter_type);
1094+
Py_VISIT(_module_state->DParameter_type);
10721095
return 0;
10731096
}
10741097

10751098
static int
10761099
_zstd_clear(PyObject *module)
10771100
{
1078-
STATE_FROM_MODULE(module);
1101+
_zstd_state* const _module_state = PyModule_GetState(module);
1102+
if (_module_state == NULL) {
1103+
return -1;
1104+
}
10791105

1080-
Py_CLEAR(MS_MEMBER(empty_bytes));
1081-
Py_CLEAR(MS_MEMBER(empty_readonly_memoryview));
1082-
Py_CLEAR(MS_MEMBER(str_read));
1083-
Py_CLEAR(MS_MEMBER(str_readinto));
1084-
Py_CLEAR(MS_MEMBER(str_write));
1085-
Py_CLEAR(MS_MEMBER(str_flush));
1106+
Py_CLEAR(_module_state->empty_bytes);
1107+
Py_CLEAR(_module_state->empty_readonly_memoryview);
1108+
Py_CLEAR(_module_state->str_read);
1109+
Py_CLEAR(_module_state->str_readinto);
1110+
Py_CLEAR(_module_state->str_write);
1111+
Py_CLEAR(_module_state->str_flush);
10861112

1087-
Py_CLEAR(MS_MEMBER(ZstdDict_type));
1088-
Py_CLEAR(MS_MEMBER(ZstdCompressor_type));
1113+
Py_CLEAR(_module_state->ZstdDict_type);
1114+
Py_CLEAR(_module_state->ZstdCompressor_type);
10891115

1090-
Py_CLEAR(MS_MEMBER(ZstdDecompressor_type));
1116+
Py_CLEAR(_module_state->ZstdDecompressor_type);
10911117

1092-
Py_CLEAR(MS_MEMBER(ZstdError));
1118+
Py_CLEAR(_module_state->ZstdError);
10931119

1094-
Py_CLEAR(MS_MEMBER(CParameter_type));
1095-
Py_CLEAR(MS_MEMBER(DParameter_type));
1120+
Py_CLEAR(_module_state->CParameter_type);
1121+
Py_CLEAR(_module_state->DParameter_type);
10961122
return 0;
10971123
}
10981124

0 commit comments

Comments
 (0)