WIP

Author: NicolasHug

src/torchcodec/_core/AVIOContextHolder.cpp

@@ -23,10 +23,10 @@ void AVIOContextHolder::createAVIOContext(
       buffer != nullptr,
       "Failed to allocate buffer of size " + std::to_string(bufferSize));
 
-  TORCH_CHECK(
-      (seek != nullptr) && ((write != nullptr) ^ (read != nullptr)),
-      "seek method must be defined, and either write or read must be defined. "
-      "But not both!")
+  //   TORCH_CHECK(
+  //       (seek != nullptr) && ((write != nullptr) ^ (read != nullptr)),
+  //       "seek method must be defined, and either write or read must be
+  //       defined. " "But not both!")
   avioContext_.reset(avioAllocContext(
       buffer,
       bufferSize,

src/torchcodec/_core/AVIOFileLikeContext.cpp

@@ -10,40 +10,20 @@
 namespace facebook::torchcodec {
 
 AVIOFileLikeContext::AVIOFileLikeContext(py::object fileLike)
-    : AVIOFileLikeContext(fileLike, false) {}
-
-std::unique_ptr<AVIOFileLikeContext> AVIOFileLikeContext::createForWriting(
-    py::object fileLike) {
-  return std::unique_ptr<AVIOFileLikeContext>(
-      new AVIOFileLikeContext(fileLike, true));
-}
-
-AVIOFileLikeContext::AVIOFileLikeContext(py::object fileLike, bool isWriteMode)
     : fileLike_{UniquePyObject(new py::object(fileLike))} {
   {
     // TODO: Is it necessary to acquire the GIL here? Is it maybe even
     // harmful? At the moment, this is only called from within a pybind
     // function, and pybind guarantees we have the GIL.
     py::gil_scoped_acquire gil;
-    if (isWriteMode) {
-      TORCH_CHECK(
-          py::hasattr(fileLike, "write"),
-          "File like object must implement a write method.");
-    } else {
-      TORCH_CHECK(
-          py::hasattr(fileLike, "read"),
-          "File like object must implement a read method.");
-    }
+    TORCH_CHECK(
+        py::hasattr(fileLike, "read"),
+        "File like object must implement a read method.");
     TORCH_CHECK(
         py::hasattr(fileLike, "seek"),
         "File like object must implement a seek method.");
   }
-
-  if (isWriteMode) {
-    createAVIOContext(nullptr, &write, &seek, &fileLike_);
-  } else {
-    createAVIOContext(&read, nullptr, &seek, &fileLike_);
-  }
+  createAVIOContext(&read, &write, &seek, &fileLike_);
 }
 
 int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
@@ -86,23 +66,6 @@ int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
   return totalNumRead == 0 ? AVERROR_EOF : totalNumRead;
 }
 
-int AVIOFileLikeContext::write(void* opaque, const uint8_t* buf, int buf_size) {
-  auto fileLike = static_cast<UniquePyObject*>(opaque);
-
-  // Note that we acquire the GIL outside of the loop. This is likely more
-  // efficient than releasing and acquiring it each loop iteration.
-  py::gil_scoped_acquire gil;
-
-  // Create a bytes object from the buffer
-  py::bytes data_bytes(reinterpret_cast<const char*>(buf), buf_size);
-
-  // Call the Python write method
-  auto bytes_written = (*fileLike)->attr("write")(data_bytes);
-
-  // Python write() should return the number of bytes written
-  return py::cast<int>(bytes_written);
-}
-
 int64_t AVIOFileLikeContext::seek(void* opaque, int64_t offset, int whence) {
   // We do not know the file size.
   if (whence == AVSEEK_SIZE) {
@@ -114,4 +77,12 @@ int64_t AVIOFileLikeContext::seek(void* opaque, int64_t offset, int whence) {
   return py::cast<int64_t>((*fileLike)->attr("seek")(offset, whence));
 }
 
+int AVIOFileLikeContext::write(void* opaque, const uint8_t* buf, int buf_size) {
+  auto fileLike = static_cast<UniquePyObject*>(opaque);
+  py::gil_scoped_acquire gil;
+  py::bytes bytes_obj(reinterpret_cast<const char*>(buf), buf_size);
+
+  return py::cast<int64_t>((*fileLike)->attr("write")(bytes_obj));
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AVIOFileLikeContext.h

@@ -15,25 +15,16 @@ namespace py = pybind11;
 
 namespace facebook::torchcodec {
 
-// Enables users to pass in a Python file-like object. We then forward all read,
-// write and seek calls back up to the methods on the Python object.
-class __attribute__((visibility("default"))) AVIOFileLikeContext
-    : public AVIOContextHolder {
+// Enables uers to pass in a Python file-like object. We then forward all read
+// and seek calls back up to the methods on the Python object.
+class AVIOFileLikeContext : public AVIOContextHolder {
  public:
-  // Constructor for reading from a file-like object
   explicit AVIOFileLikeContext(py::object fileLike);
 
-  // Constructor for writing to a file-like object
-  static std::unique_ptr<AVIOFileLikeContext> createForWriting(
-      py::object fileLike);
-
  private:
-  // Private constructor for write mode
-  AVIOFileLikeContext(py::object fileLike, bool isWriteMode);
-
   static int read(void* opaque, uint8_t* buf, int buf_size);
-  static int write(void* opaque, const uint8_t* buf, int buf_size);
   static int64_t seek(void* opaque, int64_t offset, int whence);
+  static int write(void* opaque, const uint8_t* buf, int buf_size);
 
   // Note that we dynamically allocate the Python object because we need to
   // strictly control when its destructor is called. We must hold the GIL

src/torchcodec/_core/pybind_ops.cpp

@@ -60,7 +60,7 @@ void encode_audio_to_file_like(
   audioStreamOptions.bitRate = bit_rate;
   audioStreamOptions.numChannels = num_channels;
 
-  auto avioContextHolder = AVIOFileLikeContext::createForWriting(file_like);
+  auto avioContextHolder = std::make_unique<AVIOFileLikeContext>(file_like);
 
   AudioEncoder encoder(
       samples,