Add File-like and bytes support to WavDecoder (#1461)

Author: NicolasHug

src/torchcodec/_core/AVIOContextHolder.cpp

@@ -58,4 +58,17 @@ AVIOContext* AVIOContextHolder::getAVIOContext() {
   return avioContext_.get();
 }
 
+int AVIOContextHolder::read(uint8_t*, int) {
+  STD_TORCH_CHECK(false, "read() is not supported by this AVIOContextHolder");
+}
+
+int64_t AVIOContextHolder::seek(int64_t, int) {
+  STD_TORCH_CHECK(false, "seek() is not supported by this AVIOContextHolder");
+}
+
+int64_t AVIOContextHolder::getSize() {
+  STD_TORCH_CHECK(
+      false, "getSize() is not supported by this AVIOContextHolder");
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AVIOContextHolder.h

@@ -7,6 +7,7 @@
 #pragma once
 
 #include "FFMPEGCommon.h"
+#include "StableABICompat.h"
 
 namespace facebook::torchcodec {
 
@@ -34,11 +35,18 @@ namespace facebook::torchcodec {
 //  3. A generic handle for those that just need to manage having access to an
 //     AVIOContext, but aren't necessarily concerned with how it was customized:
 //     typically, the SingleStreamDecoder.
-class AVIOContextHolder {
+class FORCE_PUBLIC_VISIBILITY AVIOContextHolder {
  public:
   virtual ~AVIOContextHolder();
   AVIOContext* getAVIOContext();
 
+  // Generic I/O primitives used by consumers that don't go through
+  // FFmpeg's AVIO layer (e.g. WavDecoder). Derived classes override
+  // the ones they support.
+  virtual int read(uint8_t* buf, int size);
+  virtual int64_t seek(int64_t offset, int whence);
+  virtual int64_t getSize();
+
  protected:
   // Make constructor protected to prevent anyone from constructing
   // an AVIOContextHolder without deriving it. (Ordinarily this would be

src/torchcodec/_core/AVIOFileContext.cpp

@@ -0,0 +1,58 @@
+// 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.
+
+#include "AVIOFileContext.h"
+
+#include <filesystem>
+#include "StableABICompat.h"
+
+namespace facebook::torchcodec {
+
+AVIOFileContext::AVIOFileContext(const std::string& path)
+    : file_(path, std::ios::binary) {
+  STD_TORCH_CHECK(file_.is_open(), "Failed to open file: ", path);
+  try {
+    fileSize_ = static_cast<int64_t>(std::filesystem::file_size(path));
+  } catch (const std::filesystem::filesystem_error& e) {
+    STD_TORCH_CHECK(
+        false, "Failed to get file size for: ", path, ". Error: ", e.what());
+  }
+}
+
+int AVIOFileContext::read(uint8_t* buf, int size) {
+  file_.read(reinterpret_cast<char*>(buf), size);
+  auto bytesRead = static_cast<int>(file_.gcount());
+  if (bytesRead == 0) {
+    return -1;
+  }
+  return bytesRead;
+}
+
+int64_t AVIOFileContext::seek(int64_t offset, int whence) {
+  std::ios_base::seekdir dir;
+  switch (whence) {
+    case SEEK_SET:
+      dir = std::ios::beg;
+      break;
+    case SEEK_CUR:
+      dir = std::ios::cur;
+      break;
+    case SEEK_END:
+      dir = std::ios::end;
+      break;
+    default:
+      return -1;
+  }
+  file_.seekg(offset, dir);
+  STD_TORCH_CHECK(!file_.fail(), "Failed to seek in file");
+  return static_cast<int64_t>(file_.tellg());
+}
+
+int64_t AVIOFileContext::getSize() {
+  return fileSize_;
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/AVIOFileContext.h

@@ -0,0 +1,32 @@
+// 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.
+
+#pragma once
+
+#include <fstream>
+#include <string>
+#include "AVIOContextHolder.h"
+
+namespace facebook::torchcodec {
+
+// For reading from a file on disk. Unlike the other AVIOContextHolder
+// subclasses, this one does NOT create an FFmpeg AVIOContext — it only
+// provides the read/seek/getSize primitives for consumers like
+// WavDecoder that do their own parsing.
+class AVIOFileContext : public AVIOContextHolder {
+ public:
+  explicit AVIOFileContext(const std::string& path);
+
+  int read(uint8_t* buf, int size) override;
+  int64_t seek(int64_t offset, int whence) override;
+  int64_t getSize() override;
+
+ private:
+  std::ifstream file_;
+  int64_t fileSize_;
+};
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/AVIOFileLikeContext.cpp

@@ -33,26 +33,54 @@ AVIOFileLikeContext::AVIOFileLikeContext(
         py::hasattr(fileLike, "seek"),
         "File like object must implement a seek method.");
   }
-  createAVIOContext(&read, &write, &seek, &fileLike_, isForWriting);
+  createAVIOContext(
+      &readCallback, &writeCallback, &seekCallback, this, isForWriting);
 }
 
-int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
-  auto fileLike = static_cast<UniquePyObject*>(opaque);
+int AVIOFileLikeContext::readCallback(
+    void* opaque,
+    uint8_t* buf,
+    int buf_size) {
+  auto self = static_cast<AVIOFileLikeContext*>(opaque);
+  int result = self->read(buf, buf_size);
+  return result < 0 ? AVERROR_EOF : result;
+}
 
-  // Note that we acquire the GIL outside of the loop. This is likely more
-  // efficient than releasing and acquiring it each loop iteration.
+int64_t
+AVIOFileLikeContext::seekCallback(void* opaque, int64_t offset, int whence) {
+  if (whence == AVSEEK_SIZE) {
+    // Size of file-like is typically unknown, since the data is potentially
+    // streaming.
+    return AVERROR(EIO);
+  }
+  auto self = static_cast<AVIOFileLikeContext*>(opaque);
+  return self->seek(offset, whence);
+}
+
+int AVIOFileLikeContext::writeCallback(
+    void* opaque,
+    const uint8_t* buf,
+    int buf_size) {
+  auto self = static_cast<AVIOFileLikeContext*>(opaque);
+  py::gil_scoped_acquire gil;
+  py::bytes bytes_obj(reinterpret_cast<const char*>(buf), buf_size);
+  return py::cast<int>(self->fileLike_->attr("write")(bytes_obj));
+}
+
+int AVIOFileLikeContext::read(uint8_t* buf, int size) {
   py::gil_scoped_acquire gil;
 
   int totalNumRead = 0;
-  while (totalNumRead < buf_size) {
-    int request = buf_size - totalNumRead;
-
-    // The Python method returns the actual bytes, which we access through the
-    // py::bytes wrapper. That wrapper, however, does not provide us access to
-    // the underlying data pointer, which we need for the memcpy below. So we
-    // convert the bytes to a string_view to get access to the data pointer.
-    // Becauase it's a view and not a copy, it should be cheap.
-    auto bytesRead = static_cast<py::bytes>((*fileLike)->attr("read")(request));
+  while (totalNumRead < size) {
+    int request = size - totalNumRead;
+
+    // The Python method returns the actual bytes, which we access through
+    // the py::bytes wrapper. That wrapper, however, does not provide us
+    // access to the underlying data pointer, which we need for the memcpy
+    // below. So we convert the bytes to a string_view to get access to
+    // the data pointer. Because it's a view and not a copy, it should be
+    // cheap.
+    auto bytesRead = static_cast<py::bytes>(fileLike_->attr("read")(request));
     auto bytesView = static_cast<std::string_view>(bytesRead);
 
     int numBytesRead = static_cast<int>(bytesView.size());
@@ -66,33 +94,26 @@ int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
         request,
         " bytes but, received ",
         numBytesRead,
-        " bytes. The given object does not conform to read protocol of file object.");
+        " bytes. The given object does not conform to read protocol "
+        "of file object.");
 
     std::memcpy(buf, bytesView.data(), numBytesRead);
     buf += numBytesRead;
     totalNumRead += numBytesRead;
   }
 
-  return totalNumRead == 0 ? AVERROR_EOF : totalNumRead;
+  return totalNumRead == 0 ? -1 : totalNumRead;
 }
 
-int64_t AVIOFileLikeContext::seek(void* opaque, int64_t offset, int whence) {
-  // We do not know the file size.
-  if (whence == AVSEEK_SIZE) {
-    return AVERROR(EIO);
-  }
-
-  auto fileLike = static_cast<UniquePyObject*>(opaque);
+int64_t AVIOFileLikeContext::seek(int64_t offset, int whence) {
   py::gil_scoped_acquire gil;
-  return py::cast<int64_t>((*fileLike)->attr("seek")(offset, 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<int>((*fileLike)->attr("write")(bytes_obj));
+int64_t AVIOFileLikeContext::getSize() {
+  // Size of file-like is typically unknown, since the data is potentially
+  // streaming.
+  return INT64_MAX;
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/AVIOFileLikeContext.h

@@ -27,10 +27,14 @@ class AVIOFileLikeContext : public AVIOContextHolder {
  public:
   explicit AVIOFileLikeContext(const py::object& fileLike, bool isForWriting);
 
+  int read(uint8_t* buf, int size) override;
+  int64_t seek(int64_t offset, int whence) override;
+  int64_t getSize() override;
+
  private:
-  static int read(void* opaque, 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);
+  static int readCallback(void* opaque, uint8_t* buf, int buf_size);
+  static int64_t seekCallback(void* opaque, int64_t offset, int whence);
+  static int writeCallback(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