@@ -554,28 +554,74 @@ static PyMethodDef _zstd_methods[] = {
554554 {0 }
555555};
556556
557-
558- static inline int
559- add_vars_to_module (PyObject * m )
557+ static int
558+ _zstd_exec (PyObject * m )
560559{
560+ #define ADD_TYPE (TYPE , SPEC ) \
561+ do { \
562+ TYPE = (PyTypeObject *)PyType_FromModuleAndSpec(m, &(SPEC), NULL); \
563+ if (TYPE == NULL) { \
564+ return -1; \
565+ } \
566+ if (PyModule_AddType(m, TYPE) < 0) { \
567+ return -1; \
568+ } \
569+ } while (0)
570+
561571#define ADD_INT_MACRO (MACRO ) \
562572 if (PyModule_AddIntConstant((m), #MACRO, (MACRO)) < 0) { \
563573 return -1; \
564574 }
565575
566- /* zstd_version_number, int */
576+ #define ADD_INT_CONST_TO_TYPE (TYPE , NAME , VALUE ) \
577+ do { \
578+ PyObject *v = PyLong_FromLong((VALUE)); \
579+ if (v == NULL || PyObject_SetAttrString((PyObject *)(TYPE), \
580+ (NAME), v) < 0) { \
581+ Py_XDECREF(v); \
582+ return -1; \
583+ } \
584+ Py_DECREF(v); \
585+ } while (0)
586+
587+ _zstd_state * const mod_state = get_zstd_state (m );
588+
589+ /* Reusable objects & variables */
590+ mod_state -> empty_bytes = PyBytes_FromStringAndSize (NULL , 0 );
591+ if (mod_state -> empty_bytes == NULL ) {
592+ return -1 ;
593+ }
594+
595+ mod_state -> CParameter_type = NULL ;
596+ mod_state -> DParameter_type = NULL ;
597+
598+ /* Create and add heap types */
599+ ADD_TYPE (mod_state -> ZstdDict_type , zstd_dict_type_spec );
600+ ADD_TYPE (mod_state -> ZstdCompressor_type , zstd_compressor_type_spec );
601+ ADD_TYPE (mod_state -> ZstdDecompressor_type , zstd_decompressor_type_spec );
602+ mod_state -> ZstdError = PyErr_NewExceptionWithDoc (
603+ "_zstd.ZstdError" ,
604+ "An error occurred in the zstd library." ,
605+ NULL , NULL );
606+ if (mod_state -> ZstdError == NULL ) {
607+ return -1 ;
608+ }
609+ if (PyModule_AddType (m , (PyTypeObject * )mod_state -> ZstdError ) < 0 ) {
610+ Py_DECREF (mod_state -> ZstdError );
611+ return -1 ;
612+ }
613+
614+ /* Add constants */
567615 if (PyModule_AddIntConstant (m , "zstd_version_number" ,
568616 ZSTD_versionNumber ()) < 0 ) {
569617 return -1 ;
570618 }
571619
572- /* zstd_version, str */
573620 if (PyModule_AddStringConstant (m , "zstd_version" ,
574621 ZSTD_versionString ()) < 0 ) {
575622 return -1 ;
576623 }
577624
578- /* ZSTD_CLEVEL_DEFAULT, int */
579625#if ZSTD_VERSION_NUMBER >= 10500
580626 if (PyModule_AddIntConstant (m , "ZSTD_CLEVEL_DEFAULT" ,
581627 ZSTD_defaultCLevel ()) < 0 ) {
@@ -585,7 +631,6 @@ add_vars_to_module(PyObject *m)
585631 ADD_INT_MACRO (ZSTD_CLEVEL_DEFAULT );
586632#endif
587633
588- /* ZSTD_DStreamOutSize, int */
589634 if (PyModule_Add (m , "ZSTD_DStreamOutSize" ,
590635 PyLong_FromSize_t (ZSTD_DStreamOutSize ())) < 0 ) {
591636 return -1 ;
@@ -618,7 +663,7 @@ add_vars_to_module(PyObject *m)
618663 /* Add zstd decompression parameters. All should also be in dp_list. */
619664 ADD_INT_MACRO (ZSTD_d_windowLogMax );
620665
621- /* ZSTD_strategy enum */
666+ /* Add ZSTD_strategy enum members */
622667 ADD_INT_MACRO (ZSTD_fast );
623668 ADD_INT_MACRO (ZSTD_dfast );
624669 ADD_INT_MACRO (ZSTD_greedy );
@@ -629,135 +674,17 @@ add_vars_to_module(PyObject *m)
629674 ADD_INT_MACRO (ZSTD_btultra );
630675 ADD_INT_MACRO (ZSTD_btultra2 );
631676
632- #undef ADD_INT_MACRO
633-
634- return 0 ;
635- }
636-
637- #define ADD_STR_TO_STATE_MACRO (STR ) \
638- do { \
639- mod_state->str_##STR = PyUnicode_FromString(#STR); \
640- if (mod_state->str_##STR == NULL) { \
641- return -1; \
642- } \
643- } while(0)
644-
645- static inline int
646- add_type_to_module (PyObject * module , const char * name ,
647- PyType_Spec * type_spec , PyTypeObject * * dest )
648- {
649- PyObject * temp = PyType_FromModuleAndSpec (module , type_spec , NULL );
650-
651- if (PyModule_AddObjectRef (module , name , temp ) < 0 ) {
652- Py_XDECREF (temp );
653- return -1 ;
654- }
655-
656- * dest = (PyTypeObject * ) temp ;
657-
658- return 0 ;
659- }
660-
661- static inline int
662- add_constant_to_type (PyTypeObject * type , const char * name , long value )
663- {
664- PyObject * temp ;
665-
666- temp = PyLong_FromLong (value );
667- if (temp == NULL ) {
668- return -1 ;
669- }
670-
671- int rc = PyObject_SetAttrString ((PyObject * ) type , name , temp );
672- Py_DECREF (temp );
673- return rc ;
674- }
675-
676- static int _zstd_exec (PyObject * module ) {
677- _zstd_state * const mod_state = get_zstd_state (module );
678-
679- /* Reusable objects & variables */
680- mod_state -> empty_bytes = PyBytes_FromStringAndSize (NULL , 0 );
681- if (mod_state -> empty_bytes == NULL ) {
682- return -1 ;
683- }
684-
685- mod_state -> empty_readonly_memoryview =
686- PyMemoryView_FromMemory ((char * )mod_state , 0 , PyBUF_READ );
687- if (mod_state -> empty_readonly_memoryview == NULL ) {
688- return -1 ;
689- }
690-
691- /* Add str to module state */
692- ADD_STR_TO_STATE_MACRO (read );
693- ADD_STR_TO_STATE_MACRO (readinto );
694- ADD_STR_TO_STATE_MACRO (write );
695- ADD_STR_TO_STATE_MACRO (flush );
696-
697- mod_state -> CParameter_type = NULL ;
698- mod_state -> DParameter_type = NULL ;
699-
700- /* Add variables to module */
701- if (add_vars_to_module (module ) < 0 ) {
702- return -1 ;
703- }
704-
705- /* ZstdError */
706- mod_state -> ZstdError = PyErr_NewExceptionWithDoc (
707- "_zstd.ZstdError" ,
708- "Call to the underlying zstd library failed." ,
709- NULL , NULL );
710- if (mod_state -> ZstdError == NULL ) {
711- return -1 ;
712- }
677+ /* Add ZSTD_EndDirective enum members to ZstdCompressor */
678+ ADD_INT_CONST_TO_TYPE (mod_state -> ZstdCompressor_type ,
679+ "CONTINUE" , ZSTD_e_continue );
680+ ADD_INT_CONST_TO_TYPE (mod_state -> ZstdCompressor_type ,
681+ "FLUSH_BLOCK" , ZSTD_e_flush );
682+ ADD_INT_CONST_TO_TYPE (mod_state -> ZstdCompressor_type ,
683+ "FLUSH_FRAME" , ZSTD_e_end );
713684
714- if (PyModule_AddObjectRef (module , "ZstdError" , mod_state -> ZstdError ) < 0 ) {
715- Py_DECREF (mod_state -> ZstdError );
716- return -1 ;
717- }
718-
719- /* ZstdDict */
720- if (add_type_to_module (module ,
721- "ZstdDict" ,
722- & zstddict_type_spec ,
723- & mod_state -> ZstdDict_type ) < 0 ) {
724- return -1 ;
725- }
726-
727- // ZstdCompressor
728- if (add_type_to_module (module ,
729- "ZstdCompressor" ,
730- & zstdcompressor_type_spec ,
731- & mod_state -> ZstdCompressor_type ) < 0 ) {
732- return -1 ;
733- }
734-
735- // Add EndDirective enum to ZstdCompressor
736- if (add_constant_to_type (mod_state -> ZstdCompressor_type ,
737- "CONTINUE" ,
738- ZSTD_e_continue ) < 0 ) {
739- return -1 ;
740- }
741-
742- if (add_constant_to_type (mod_state -> ZstdCompressor_type ,
743- "FLUSH_BLOCK" ,
744- ZSTD_e_flush ) < 0 ) {
745- return -1 ;
746- }
747-
748- if (add_constant_to_type (mod_state -> ZstdCompressor_type ,
749- "FLUSH_FRAME" ,
750- ZSTD_e_end ) < 0 ) {
751- return -1 ;
752- }
753-
754- // ZstdDecompressor
755- if (add_type_to_module (module ,
756- "ZstdDecompressor" ,
757- & zstddecompressor_type_spec ,
758- & mod_state -> ZstdDecompressor_type ) < 0 ) {
759- return -1 ;
760- }
685+ #undef ADD_TYPE
686+ #undef ADD_INT_MACRO
687+ #undef ADD_ZSTD_COMPRESSOR_INT_CONST
761688
762689 return 0 ;
763690}
@@ -768,11 +695,6 @@ _zstd_traverse(PyObject *module, visitproc visit, void *arg)
768695 _zstd_state * const mod_state = get_zstd_state (module );
769696
770697 Py_VISIT (mod_state -> empty_bytes );
771- Py_VISIT (mod_state -> empty_readonly_memoryview );
772- Py_VISIT (mod_state -> str_read );
773- Py_VISIT (mod_state -> str_readinto );
774- Py_VISIT (mod_state -> str_write );
775- Py_VISIT (mod_state -> str_flush );
776698
777699 Py_VISIT (mod_state -> ZstdDict_type );
778700 Py_VISIT (mod_state -> ZstdCompressor_type );
@@ -792,11 +714,6 @@ _zstd_clear(PyObject *module)
792714 _zstd_state * const mod_state = get_zstd_state (module );
793715
794716 Py_CLEAR (mod_state -> empty_bytes );
795- Py_CLEAR (mod_state -> empty_readonly_memoryview );
796- Py_CLEAR (mod_state -> str_read );
797- Py_CLEAR (mod_state -> str_readinto );
798- Py_CLEAR (mod_state -> str_write );
799- Py_CLEAR (mod_state -> str_flush );
800717
801718 Py_CLEAR (mod_state -> ZstdDict_type );
802719 Py_CLEAR (mod_state -> ZstdCompressor_type );
0 commit comments