use module-bound exception

PiperOrigin-RevId: 811366325

Author: eustas

python/_brotli.c

@@ -18,7 +18,8 @@
 #define PY_GET_TYPE(Obj) ((Obj)->ob_type)
 #endif
 
-static PyObject* BrotliError;
+static const char kErrorAttr[] = "error";
+static const char kModuleAttr[] = "_module";
 
 static const char kInvalidBufferError[] =
     "brotli: data must be a C-contiguous buffer";
@@ -195,8 +196,26 @@ PyDoc_STRVAR(brotli_doc,
 "Implementation module for the Brotli library.");
 /* clang-format on */
 
-static void set_brotli_exception(const char* msg) {
-  PyErr_SetString(BrotliError, msg);
+static void set_brotli_exception(PyObject* context, const char* msg) {
+  assert(context != NULL);
+  PyObject* module = NULL;
+  PyObject* type = NULL;
+  PyObject* error = NULL;
+  if (PyModule_Check(context)) {
+    module = context;
+  } else if (PyType_Check(context)) {
+    type = context;
+  } else {
+    type = (PyObject*)PY_GET_TYPE(context);
+    assert(type != NULL);
+  }
+  if (type) {
+    module = PyObject_GetAttrString(type, kModuleAttr);
+  }
+  assert(module != NULL);
+  error = PyObject_GetAttrString(module, kErrorAttr);
+  assert(error != NULL);
+  PyErr_SetString(error, msg);
 }
 
 /*
@@ -354,7 +373,7 @@ static PyObject* brotli_Compressor_new(PyTypeObject* type, PyObject* args,
   self->processing = 0;
   self->enc = BrotliEncoderCreateInstance(0, 0, 0);
   if (self->enc == NULL) {
-    set_brotli_exception(kCompressCreateError);
+    set_brotli_exception((PyObject*)type, kCompressCreateError);
     PY_GET_TYPE(self)->tp_free((PyObject*)self);
     return NULL;
   }
@@ -386,30 +405,30 @@ static int brotli_Compressor_init(PyBrotli_Compressor* self, PyObject* args,
   if ((mode == 0) || (mode == 1) || (mode == 2)) {
     BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_MODE, (uint32_t)mode);
   } else {
-    set_brotli_exception(kInvalidModeError);
+    set_brotli_exception((PyObject*)self, kInvalidModeError);
     self->healthy = 0;
     return -1;
   }
   if (quality <= 11) {
     BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_QUALITY,
                               (uint32_t)quality);
   } else {
-    set_brotli_exception(kInvalidQualityError);
+    set_brotli_exception((PyObject*)self, kInvalidQualityError);
     self->healthy = 0;
     return -1;
   }
   if ((10 <= lgwin) && (lgwin <= 24)) {
     BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_LGWIN, (uint32_t)lgwin);
   } else {
-    set_brotli_exception(kInvalidLgwinError);
+    set_brotli_exception((PyObject*)self, kInvalidLgwinError);
     self->healthy = 0;
     return -1;
   }
   if ((lgblock == 0) || ((16 <= lgblock) && (lgblock <= 24))) {
     BrotliEncoderSetParameter(self->enc, BROTLI_PARAM_LGBLOCK,
                               (uint32_t)lgblock);
   } else {
-    set_brotli_exception(kInvalidLgblockError);
+    set_brotli_exception((PyObject*)self, kInvalidLgblockError);
     self->healthy = 0;
     return -1;
   }
@@ -471,7 +490,7 @@ static PyObject* compress_stream(PyBrotli_Compressor* self,
     ret = Buffer_Finish(&buffer);
     if (ret == NULL) oom = 1;
   } else {  /* Not ok */
-    set_brotli_exception(kCompressError);
+    set_brotli_exception((PyObject*)self, kCompressError);
   }
 
 error:
@@ -495,11 +514,11 @@ static PyObject* brotli_Compressor_process(PyBrotli_Compressor* self,
   Py_buffer input;
 
   if (self->healthy == 0) {
-    set_brotli_exception(kCompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kCompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kCompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kCompressConcurrentError);
     return NULL;
   }
 
@@ -522,11 +541,11 @@ static PyObject* brotli_Compressor_flush(PyBrotli_Compressor* self) {
   PyObject* ret = NULL;
 
   if (self->healthy == 0) {
-    set_brotli_exception(kCompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kCompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kCompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kCompressConcurrentError);
     return NULL;
   }
 
@@ -540,11 +559,11 @@ static PyObject* brotli_Compressor_finish(PyBrotli_Compressor* self) {
   PyObject* ret = NULL;
 
   if (self->healthy == 0) {
-    set_brotli_exception(kCompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kCompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kCompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kCompressConcurrentError);
     return NULL;
   }
 
@@ -578,7 +597,7 @@ static PyObject* brotli_Decompressor_new(PyTypeObject* type, PyObject* args,
 
   self->dec = BrotliDecoderCreateInstance(0, 0, 0);
   if (self->dec == NULL) {
-    set_brotli_exception(kDecompressCreateError);
+    set_brotli_exception((PyObject*)type, kDecompressCreateError);
     PY_GET_TYPE(self)->tp_free((PyObject*)self);
     return NULL;
   }
@@ -633,11 +652,11 @@ static PyObject* brotli_Decompressor_process(PyBrotli_Decompressor* self,
   int oom = 0;
 
   if (self->healthy == 0) {
-    set_brotli_exception(kDecompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kDecompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kDecompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kDecompressConcurrentError);
     return NULL;
   }
 
@@ -654,7 +673,7 @@ static PyObject* brotli_Decompressor_process(PyBrotli_Decompressor* self,
 
   if (self->unconsumed_data_length > 0) {
     if (input.len > 0) {
-      set_brotli_exception(kDecompressSinkError);
+      set_brotli_exception((PyObject*)self, kDecompressSinkError);
       goto finally;
     }
     next_in = self->unconsumed_data;
@@ -693,7 +712,7 @@ static PyObject* brotli_Decompressor_process(PyBrotli_Decompressor* self,
   if (oom) {
     goto finally;
   } else if (result == BROTLI_DECODER_RESULT_ERROR) {
-    set_brotli_exception(kDecompressError);
+    set_brotli_exception((PyObject*)self, kDecompressError);
     goto finally;
   }
 
@@ -712,7 +731,7 @@ static PyObject* brotli_Decompressor_process(PyBrotli_Decompressor* self,
 
   if ((result == BROTLI_DECODER_RESULT_SUCCESS) && (avail_in > 0)) {
     /* TODO(eustas): Add API to ignore / fetch unused "tail"? */
-    set_brotli_exception(kDecompressError);
+    set_brotli_exception((PyObject*)self, kDecompressError);
     goto finally;
   }
 
@@ -742,11 +761,11 @@ static PyObject* brotli_Decompressor_process(PyBrotli_Decompressor* self,
 
 static PyObject* brotli_Decompressor_is_finished(PyBrotli_Decompressor* self) {
   if (self->healthy == 0) {
-    set_brotli_exception(kDecompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kDecompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kDecompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kDecompressConcurrentError);
     return NULL;
   }
   if (BrotliDecoderIsFinished(self->dec)) {
@@ -759,11 +778,11 @@ static PyObject* brotli_Decompressor_is_finished(PyBrotli_Decompressor* self) {
 static PyObject* brotli_Decompressor_can_accept_more_data(
     PyBrotli_Decompressor* self) {
   if (self->healthy == 0) {
-    set_brotli_exception(kDecompressUnhealthyError);
+    set_brotli_exception((PyObject*)self, kDecompressUnhealthyError);
     return NULL;
   }
   if (self->processing != 0) {
-    set_brotli_exception(kDecompressConcurrentError);
+    set_brotli_exception((PyObject*)self, kDecompressConcurrentError);
     return NULL;
   }
   if (self->unconsumed_data_length > 0) {
@@ -832,7 +851,7 @@ static PyObject* brotli_decompress(PyObject* self, PyObject* args,
   if (oom) {
     goto finally;
   } else if (result != BROTLI_DECODER_RESULT_SUCCESS || available_in > 0) {
-    set_brotli_exception(kDecompressError);
+    set_brotli_exception((PyObject*)self, kDecompressError);
     goto finally;
   }
 
@@ -1032,16 +1051,15 @@ PyMODINIT_FUNC INIT_BROTLI(void) {
   PyObject* error_type = NULL;
   PyObject* compressor_type = NULL;
   PyObject* decompressor_type = NULL;
+  PyObject* obj = NULL;  /* Dangling ref, don't deref this one. */
 
   if (m == NULL) goto error;
 
   error_type = PyErr_NewExceptionWithDoc((char*)"brotli.error",
                                          brotli_error_doc, NULL, NULL);
   if (error_type == NULL) goto error;
 
-  if (RegisterObject(m, "error", error_type) < 0) goto error;
-  /* Assumption: pointer is used only while module is alive and well. */
-  BrotliError = error_type;
+  if (RegisterObject(m, kErrorAttr, error_type) < 0) goto error;
   error_type = NULL;
 
 #if PY_MAJOR_VERSION >= 3
@@ -1056,12 +1074,16 @@ PyMODINIT_FUNC INIT_BROTLI(void) {
   if (compressor_type == NULL) goto error;
   if (PyType_Ready((PyTypeObject*)compressor_type) < 0) goto error;
   if (RegisterObject(m, "Compressor", compressor_type) < 0) goto error;
+  obj = compressor_type;
   compressor_type = NULL;
+  if (PyObject_SetAttrString(obj, kModuleAttr, m) < 0) goto error;
 
   if (decompressor_type == NULL) goto error;
   if (PyType_Ready((PyTypeObject*)decompressor_type) < 0) goto error;
   if (RegisterObject(m, "Decompressor", decompressor_type) < 0) goto error;
+  obj = decompressor_type;
   decompressor_type = NULL;
+  if (PyObject_SetAttrString(obj, kModuleAttr, m) < 0) goto error;
 
   PyModule_AddIntConstant(m, "MODE_GENERIC", (int)BROTLI_MODE_GENERIC);
   PyModule_AddIntConstant(m, "MODE_TEXT", (int)BROTLI_MODE_TEXT);