Store original data ref in reference audio/video frames (#1176)

Previously, `create_reference_audio_frame` and
`create_reference_video_frame` created views to the underlying buffer
without keeping a reference to the original data.
This was dangerous because if the original array was garbage collected
while the frame object still existed,
it would lead to undefined behavior and potential segmentation faults.

This commit resolves this by:

- Adding dynamic attributes to Frame classes: Modified the nanobind
class registration for `AudioFrames` and `VideoFrames` in `frames.cpp`
to include `nb::dynamic_attr()`, which allows Python to add custom
attributes to these C++ objects.

- Storing reference to original data: Updated the Python implementations
of `create_reference_audio_frame` and `create_reference_video_frame` to
store the original array in `frame._array_ref` after creating the frame
object. This prevents the original data from being garbage collected as
long as the frame object is alive.

Author: mthrok

src/spdl/io/_core.py

@@ -1143,12 +1143,12 @@ def create_reference_audio_frame(
     This function should be used when the media data processed in Python should
     be further processed by filter graph, and/or encoded.
 
-    .. attention::
+    .. note::
 
        The resulting frame object references the memory region owned by the input
-       array, but it does not own the reference to the original array.
-
-       Make sure that the array object is alive until the frame object is consumed.
+       array and keeps a reference to the original array to prevent it from being
+       garbage collected. The original data will remain alive as long as the frame
+       object is alive.
 
     Args:
         array: 2D array or tensor.
@@ -1177,12 +1177,14 @@ def create_reference_audio_frame(
     Returns:
        Frames object that references the memory region of the input data.
     """
-    return _libspdl.create_reference_audio_frame(
+    frame = _libspdl.create_reference_audio_frame(
         array=array,
         sample_fmt=sample_fmt,
         sample_rate=sample_rate,
         pts=pts,
     )
+    frame._array_ref = array
+    return frame
 
 
 def create_reference_video_frame(
@@ -1193,12 +1195,12 @@ def create_reference_video_frame(
     This function should be used when the media data processed in Python should
     be further processed by filter graph, and/or encoded.
 
-    .. attention::
+    .. note::
 
        The resulting frame object references the memory region owned by the input
-       array, but it does not own the reference to the original array.
-
-       Make sure that the array object is alive until the frame object is consumed.
+       array and keeps a reference to the original array to prevent it from being
+       garbage collected. The original data will remain alive as long as the frame
+       object is alive.
 
     Args:
         array: 3D or 4D array or tensor.
@@ -1221,12 +1223,14 @@ def create_reference_video_frame(
     Returns:
        Frames object that references the memory region of the input data.
     """
-    return _libspdl.create_reference_video_frame(
+    frame = _libspdl.create_reference_video_frame(
         array=array,
         pix_fmt=pix_fmt,
         frame_rate=frame_rate,
         pts=pts,
     )
+    frame._array_ref = array
+    return frame
 
 
 ################################################################################

src/spdl/io/lib/core/frames.cpp

@@ -66,7 +66,8 @@ void register_frames(nb::module_& m) {
       m,
       "AudioFrames",
       "Audio frames.\n\n"
-      "See :doc:`/io/packets_frames_concepts` for information about the Frames base concept.")
+      "See :doc:`/io/packets_frames_concepts` for information about the Frames base concept.",
+      nb::dynamic_attr())
       .def_prop_ro(
           "num_frames",
           [](AudioFrames& self) {
@@ -133,7 +134,8 @@ void register_frames(nb::module_& m) {
       m,
       "VideoFrames",
       "Video frames.\n\n"
-      "See :doc:`/io/packets_frames_concepts` for information about the Frames base concept.")
+      "See :doc:`/io/packets_frames_concepts` for information about the Frames base concept.",
+      nb::dynamic_attr())
       .def_prop_ro(
           "num_frames",
           [](VideoFrames& self) {

tests/io/reference_frames_test.py

@@ -0,0 +1,166 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import gc
+import unittest
+import weakref
+
+import numpy as np
+import spdl.io
+
+
+class TestReferenceAudioFrame(unittest.TestCase):
+    def test_audio_frame_keeps_reference_alive(self) -> None:
+        """Test that AudioFrame keeps the original array alive even after deletion."""
+        # Setup: Create an array and a weak reference to it
+        array = np.random.randint(0, 255, size=(100, 2), dtype=np.uint8)
+        weak_ref = weakref.ref(array)
+
+        # Execute: Create a reference frame
+        frame = spdl.io.create_reference_audio_frame(
+            array=array,
+            sample_fmt="u8",
+            sample_rate=44100,
+            pts=0,
+        )
+
+        # Assert: The array should still be alive (referenced by both array and frame)
+        self.assertIsNotNone(weak_ref())
+
+        # Execute: Delete the original array variable
+        del array
+        gc.collect()
+
+        # Assert: The array should still be alive (referenced by frame)
+        self.assertIsNotNone(weak_ref())
+
+        # Execute: Delete the frame
+        del frame
+        gc.collect()
+
+        # Assert: Now the array should be garbage collected
+        self.assertIsNone(weak_ref())
+
+    def test_audio_frame_can_be_converted(self) -> None:
+        """Test that reference audio frame can be converted to buffer."""
+        array = np.random.randint(0, 255, size=(100, 2), dtype=np.uint8)
+
+        frame = spdl.io.create_reference_audio_frame(
+            array=array,
+            sample_fmt="u8",
+            sample_rate=44100,
+            pts=0,
+        )
+        buffer = spdl.io.convert_frames(frame)
+
+        buffer_array = spdl.io.to_numpy(buffer)
+        self.assertEqual(buffer_array.shape, (100, 2))
+        np.testing.assert_array_equal(buffer_array, array)
+
+    def test_audio_frame_planar_format(self) -> None:
+        """Test that planar format audio frames work correctly."""
+        # Setup: Create planar format data (channels first)
+        array = np.random.randint(0, 255, size=(2, 100), dtype=np.uint8)
+
+        # Execute: Create reference frame with planar format
+        frame = spdl.io.create_reference_audio_frame(
+            array=array,
+            sample_fmt="u8p",
+            sample_rate=44100,
+            pts=0,
+        )
+
+        # Assert: Frame should have correct properties
+        self.assertEqual(frame.num_frames, 100)
+        self.assertEqual(frame.num_channels, 2)
+        self.assertEqual(frame.sample_rate, 44100)
+
+
+class TestReferenceVideoFrame(unittest.TestCase):
+    def test_video_frame_keeps_reference_alive(self) -> None:
+        """Test that VideoFrame keeps the original array alive even after deletion."""
+        # Setup: Create an array and a weak reference to it
+        array = np.random.randint(0, 255, size=(5, 128, 128, 3), dtype=np.uint8)
+        weak_ref = weakref.ref(array)
+
+        # Execute: Create a reference frame
+        frame = spdl.io.create_reference_video_frame(
+            array=array,
+            pix_fmt="rgb24",
+            frame_rate=(30, 1),
+            pts=0,
+        )
+
+        # Assert: The array should still be alive (referenced by both array and frame)
+        self.assertIsNotNone(weak_ref())
+
+        # Execute: Delete the original array variable
+        del array
+        gc.collect()
+
+        # Assert: The array should still be alive (referenced by frame)
+        self.assertIsNotNone(weak_ref())
+
+        # Execute: Delete the frame
+        del frame
+        gc.collect()
+
+        # Assert: Now the array should be garbage collected
+        self.assertIsNone(weak_ref())
+
+    def test_video_frame_can_be_converted(self) -> None:
+        """Test that reference video frame can be converted to buffer."""
+        array = np.random.randint(0, 255, size=(5, 128, 128, 3), dtype=np.uint8)
+
+        frame = spdl.io.create_reference_video_frame(
+            array=array,
+            pix_fmt="rgb24",
+            frame_rate=(30, 1),
+            pts=0,
+        )
+        buffer = spdl.io.convert_frames(frame)
+
+        buffer_array = spdl.io.to_numpy(buffer)
+        self.assertEqual(buffer_array.shape, (5, 128, 128, 3))
+        np.testing.assert_array_equal(buffer_array, array)
+
+    def test_video_frame_grayscale(self) -> None:
+        """Test that grayscale video frames work correctly."""
+        # Setup: Create grayscale data
+        array = np.random.randint(0, 255, size=(5, 128, 128), dtype=np.uint8)
+
+        # Execute: Create reference frame with grayscale format
+        frame = spdl.io.create_reference_video_frame(
+            array=array,
+            pix_fmt="gray8",
+            frame_rate=(30, 1),
+            pts=0,
+        )
+
+        # Assert: Frame should have correct properties
+        self.assertEqual(frame.num_frames, 5)
+        self.assertEqual(frame.width, 128)
+        self.assertEqual(frame.height, 128)
+
+    def test_video_frame_yuv_format(self) -> None:
+        """Test that YUV planar format video frames work correctly."""
+        # Setup: Create YUV planar format data (channels first)
+        array = np.random.randint(0, 255, size=(5, 3, 128, 128), dtype=np.uint8)
+
+        # Execute: Create reference frame with YUV planar format
+        frame = spdl.io.create_reference_video_frame(
+            array=array,
+            pix_fmt="yuv444p",
+            frame_rate=(30, 1),
+            pts=0,
+        )
+
+        # Assert: Frame should have correct properties
+        self.assertEqual(frame.num_frames, 5)
+        self.assertEqual(frame.width, 128)
+        self.assertEqual(frame.height, 128)