Add constants to ZstdCompressor in C

Author: AA-Turner

Lib/compression/zstd/__init__.py

@@ -40,11 +40,6 @@
 COMPRESSION_LEVEL_DEFAULT = _zstd.ZSTD_CLEVEL_DEFAULT
 """The default compression level for Zstandard, currently '3'."""
 
-# Set ZSTD_EndDirective constants on ZstdCompressor
-ZstdCompressor.CONTINUE = _zstd.ZSTD_e_continue
-ZstdCompressor.FLUSH_BLOCK = _zstd.ZSTD_e_flush
-ZstdCompressor.FLUSH_FRAME = _zstd.ZSTD_e_end
-
 
 class FrameInfo:
     """Information about a Zstandard frame."""

Lib/compression/zstd/_zstdfile.py

@@ -1,7 +1,7 @@
 import io
 from os import PathLike
 from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdError,
-                   ZSTD_DStreamOutSize, ZSTD_e_end, ZSTD_e_flush)
+                   ZSTD_DStreamOutSize)
 from compression._common import _streams
 
 __all__ = ("ZstdFile", "open")
@@ -28,8 +28,8 @@ class ZstdFile(_streams.BaseStream):
     bytes, and may only be written to objects that support the Buffer Protocol.
     """
 
-    FLUSH_BLOCK = ZSTD_e_flush  # ZstdCompressor.FLUSH_BLOCK
-    FLUSH_FRAME = ZSTD_e_end  # ZstdCompressor.FLUSH_FRAME
+    FLUSH_BLOCK = ZstdCompressor.FLUSH_BLOCK
+    FLUSH_FRAME = ZstdCompressor.FLUSH_FRAME
 
     def __init__(self, file, /, mode="r", *,
                  level=None, options=None, zstd_dict=None):

Modules/_zstd/_zstdmodule.c

@@ -573,6 +573,17 @@ do {                                                                         \
         return -1;                                                           \
     }
 
+#define ADD_INT_CONST_TO_TYPE(TYPE, NAME, VALUE)                             \
+do {                                                                         \
+    PyObject *v = PyLong_FromLong((VALUE));                                  \
+    if (v == NULL || PyObject_SetAttrString((PyObject *)(TYPE),              \
+                                            (NAME), v) < 0) {                \
+        Py_XDECREF(v);                                                       \
+        return -1;                                                           \
+    }                                                                        \
+    Py_DECREF(v);                                                            \
+} while (0)
+
     _zstd_state* const mod_state = get_zstd_state(m);
 
     /* Reusable objects & variables */
@@ -663,13 +674,17 @@ do {                                                                         \
     ADD_INT_MACRO(ZSTD_btultra);
     ADD_INT_MACRO(ZSTD_btultra2);
 
-    /* Add ZSTD_EndDirective enum members */
-    ADD_INT_MACRO(ZSTD_e_continue);
-    ADD_INT_MACRO(ZSTD_e_flush);
-    ADD_INT_MACRO(ZSTD_e_end);
+    /* Add ZSTD_EndDirective enum members to ZstdCompressor */
+    ADD_INT_CONST_TO_TYPE(mod_state->ZstdCompressor_type,
+                          "CONTINUE", ZSTD_e_continue);
+    ADD_INT_CONST_TO_TYPE(mod_state->ZstdCompressor_type,
+                          "FLUSH_BLOCK", ZSTD_e_flush);
+    ADD_INT_CONST_TO_TYPE(mod_state->ZstdCompressor_type,
+                          "FLUSH_FRAME", ZSTD_e_end);
 
 #undef ADD_TYPE
 #undef ADD_INT_MACRO
+#undef ADD_ZSTD_COMPRESSOR_INT_CONST
 
     return 0;
 }