0.1.10 - fix decompress data refcount issue

K0lb3 · Dec 7, 2024 · 4ba915f · 4ba915f
1 parent a0d704a
commit 4ba915f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/astc_encoder/__init__.py b/astc_encoder/__init__.py
@@ -6,7 +6,7 @@
 The exception is aarch64, which uses the neon encoder by default.
 """
 
-__version__ = "0.1.9"
+__version__ = "0.1.10"
 
 from .enum import (
     ASTCConfigFlags as ASTCConfigFlags,

diff --git a/src/pybind.cpp b/src/pybind.cpp
@@ -673,7 +673,6 @@ PyObject *ASTCContext_method_decompress(ASTContextT *self, PyObject *args, PyObj
     }
 
     // create a python bytes object from the decompressed data
-    Py_IncRef(py_image_data);
     Py_DecRef(py_image->data);
     py_image->data = py_image_data;
 

diff --git a/tests/test_main.py b/tests/test_main.py
@@ -1,9 +1,10 @@
+import gc
 import os
+import sys
 from typing import Tuple
-import gc
-import psutil
 
 import imagehash
+import psutil
 from PIL import Image
 
 import astc_encoder
@@ -145,6 +146,8 @@ def assert_memory_cleaned():
         # a leak in C/C++ would most likely not align with this
         assert (current_memory - initial_memory) % 4096 == 0
 
+    total_obj_count = len(gc.get_objects())
+
     # create and destroy all object types
     astc_encoder.ASTCSwizzle()
     assert_memory_cleaned()
@@ -169,17 +172,28 @@ def assert_memory_cleaned():
     # do a full round of compression/decompression
     def full_run():
         swizzle = astc_encoder.ASTCSwizzle()
+        astc_image_data = b"\00" * 64
+        astc_image_data_refs = sys.getrefcount(astc_image_data)
         astc_image = astc_encoder.ASTCImage(
-            astc_encoder.ASTCType.U8, 4, 4, data=b"\00" * 64
+            astc_encoder.ASTCType.U8, 4, 4, 1, astc_image_data
         )
         config = astc_encoder.ASTCConfig(astc_encoder.ASTCProfile.LDR_SRGB, 4, 4)
         context = astc_encoder.ASTCContext(config)
         comp = context.compress(astc_image, swizzle)
         astc_image_new = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, 4, 4)
         decomp = context.decompress(comp, astc_image_new, swizzle)
 
+        assert sys.getrefcount(astc_image) == 2
+        assert sys.getrefcount(astc_image.data) == astc_image_data_refs + 2
+        del astc_image
+        assert sys.getrefcount(astc_image_data) == astc_image_data_refs
+
+        assert sys.getrefcount(astc_image_new) == 3  # inp, astc_image_new, decomp
+        assert sys.getrefcount(astc_image_new.data) == 3
+
     full_run()
     assert_memory_cleaned()
+    assert total_obj_count >= len(gc.get_objects())
 
 
 if __name__ == "__main__":