Add test, reader and writer

AIDASoft · Jan 21, 2024 · 5642b3a · 5642b3a
1 parent c04312c
commit 5642b3a
Show file tree

Hide file tree

Showing 8 changed files with 473 additions and 0 deletions.
diff --git a/include/podio/Reader.h b/include/podio/Reader.h
@@ -0,0 +1,96 @@
+#ifndef PODIO_READER_H
+#define PODIO_READER_H
+
+#include "podio/Frame.h"
+#include "podio/podioVersion.h"
+
+namespace podio {
+
+class Reader {
+public:
+  struct ReaderConcept {
+    virtual ~ReaderConcept() = default;
+
+    virtual podio::Frame readNextFrame(const std::string& name) = 0;
+    virtual podio::Frame readFrame(const std::string& name, size_t index) = 0;
+    virtual size_t getEntries(const std::string& name) = 0;
+    virtual podio::Frame readNextEvent() = 0;
+    virtual podio::Frame readEvent(size_t index) = 0;
+    virtual size_t getEvents() = 0;
+    virtual podio::version::Version currentFileVersion() const = 0;
+  };
+
+  template <typename T>
+  struct ReaderModel : public ReaderConcept {
+    ReaderModel(T* reader) : m_reader(reader) {
+    }
+    ReaderModel(const ReaderModel&) = delete;
+    ReaderModel& operator=(const ReaderModel&) = delete;
+
+    podio::Frame readNextFrame(const std::string& name) override {
+      auto maybeFrame = m_reader->readNextEntry(name);
+      if (maybeFrame) {
+        return std::move(maybeFrame);
+      }
+      throw std::runtime_error("Could not read frame (reading beyond bounds?)");
+    }
+    podio::Frame readNextEvent() override {
+      return readNextFrame(podio::Category::Event);
+    }
+
+    podio::Frame readFrame(const std::string& name, size_t index) override {
+      auto maybeFrame = m_reader->readEntry(name, index);
+      if (maybeFrame) {
+        return std::move(maybeFrame);
+      }
+      throw std::runtime_error("Could not read frame (reading beyond bounds?)");
+    }
+    podio::Frame readEvent(size_t index) override {
+      return readFrame(podio::Category::Event, index);
+    }
+    size_t getEntries(const std::string& name) override {
+      return m_reader->getEntries(name);
+    }
+    size_t getEvents() override {
+      return getEntries(podio::Category::Event);
+    }
+    podio::version::Version currentFileVersion() const override {
+      return m_reader->currentFileVersion();
+    }
+    T* m_reader{nullptr};
+  };
+
+  std::unique_ptr<ReaderConcept> m_self{nullptr};
+
+  template <typename T>
+  Reader(std::unique_ptr<T>);
+
+  podio::Frame readNextFrame(const std::string& name) {
+    return m_self->readNextFrame(name);
+  }
+  podio::Frame readNextEvent() {
+    return readNextFrame(podio::Category::Event);
+  }
+  podio::Frame readFrame(const std::string& name, size_t index) {
+    return m_self->readFrame(name, index);
+  }
+  podio::Frame readEvent(size_t index) {
+    return readFrame(podio::Category::Event, index);
+  }
+  size_t getEntries(const std::string& name) {
+    return m_self->getEntries(name);
+  }
+  size_t getEvents() {
+    return getEntries(podio::Category::Event);
+  }
+  podio::version::Version currentFileVersion() const {
+    return m_self->currentFileVersion();
+  }
+};
+
+std::unique_ptr<Reader> makeReader(const std::string& filename);
+std::unique_ptr<Reader> makeReader(const std::vector<std::string>& filename);
+
+} // namespace podio
+
+#endif // PODIO_READER_H
diff --git a/include/podio/Writer.h b/include/podio/Writer.h
@@ -0,0 +1,67 @@
+#ifndef PODIO_WRITER_H
+#define PODIO_WRITER_H
+
+#include "podio/Frame.h"
+#include "podio/podioVersion.h"
+
+namespace podio {
+
+class Writer {
+public:
+  struct WriterConcept {
+    virtual ~WriterConcept() = default;
+
+    virtual void writeFrame(const podio::Frame& frame, const std::string& category) = 0;
+    virtual void writeFrame(const podio::Frame& frame, const std::string& category,
+                            const std::vector<std::string>& collections) = 0;
+    virtual void writeEvent(const podio::Frame& frame) = 0;
+    virtual void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) = 0;
+  };
+
+  template <typename T>
+  struct WriterModel : public WriterConcept {
+    WriterModel(T* writer) : m_writer(writer) {
+    }
+    WriterModel(const WriterModel&) = delete;
+    WriterModel& operator=(const WriterModel&) = delete;
+
+    void writeFrame(const podio::Frame& frame, const std::string& category) override {
+      return m_writer->writeFrame(frame, category);
+    }
+    void writeFrame(const podio::Frame& frame, const std::string& category,
+                    const std::vector<std::string>& collections) override {
+      return m_writer->writeFrame(frame, category, collections);
+    }
+    void writeEvent(const podio::Frame& frame) override {
+      return writeFrame(frame, podio::Category::Event);
+    }
+    void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) override {
+      return writeFrame(frame, podio::Category::Event, collections);
+    }
+    std::unique_ptr<T> m_writer{nullptr};
+  };
+
+  std::unique_ptr<WriterConcept> m_self{nullptr};
+
+  template <typename T>
+  Writer(std::unique_ptr<T>);
+
+  void writeFrame(const podio::Frame& frame, const std::string& category) {
+    return m_self->writeFrame(frame, category);
+  }
+  void writeFrame(const podio::Frame& frame, const std::string& category, const std::vector<std::string>& collections) {
+    return m_self->writeFrame(frame, category, collections);
+  }
+  void writeEvent(const podio::Frame& frame) {
+    return writeFrame(frame, podio::Category::Event);
+  }
+  void writeEvent(const podio::Frame& frame, const std::vector<std::string>& collections) {
+    return writeFrame(frame, podio::Category::Event, collections);
+  }
+};
+
+std::unique_ptr<Writer> makeWriter(const std::string& filename, const std::string& type = "default");
+
+} // namespace podio
+
+#endif // PODIO_WRITER_H
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -80,6 +80,8 @@ SET(root_sources
   ROOTFrameWriter.cc
   ROOTFrameReader.cc
   ROOTLegacyReader.cc
+  Reader.cc
+  Writer.cc
 )
 if(ENABLE_RNTUPLE)
   list(APPEND root_sources
@@ -92,6 +94,8 @@ SET(root_headers
   ${PROJECT_SOURCE_DIR}/include/podio/ROOTFrameReader.h
   ${PROJECT_SOURCE_DIR}/include/podio/ROOTLegacyReader.h
   ${PROJECT_SOURCE_DIR}/include/podio/ROOTFrameWriter.h
+  ${PROJECT_SOURCE_DIR}/include/podio/Reader.h
+  ${PROJECT_SOURCE_DIR}/include/podio/Writer.h
   )
 if(ENABLE_RNTUPLE)
   list(APPEND root_headers

diff --git a/src/Reader.cc b/src/Reader.cc
@@ -0,0 +1,77 @@
+#include "podio/Reader.h"
+
+#include "podio/ROOTFrameReader.h"
+#ifdef PODIO_ENABLE_RNTUPLE
+  #include "podio/ROOTRNTupleReader.h"
+#endif
+#ifdef PODIO_ENABLE_SIO
+  #include "podio/SIOFrameReader.h"
+#endif
+
+#include "TFile.h"
+#include "TKey.h"
+#include <memory>
+
+namespace podio {
+
+template <typename T>
+Reader::Reader(std::unique_ptr<T> reader) : m_self(std::make_unique<ReaderModel<T>>(reader.release())) {
+}
+
+std::unique_ptr<Reader> makeReader(const std::string& filename) {
+  return makeReader(std::vector<std::string>{filename});
+}
+
+std::unique_ptr<Reader> makeReader(const std::vector<std::string>& filenames) {
+
+  auto suffix = filenames[0].substr(filenames[0].find_last_of(".") + 1);
+  for (size_t i = 1; i < filenames.size(); ++i) {
+    if (filenames[i].substr(filenames[i].find_last_of(".") + 1) != suffix) {
+      std::cout << "ERROR: All files must have the same extension" << std::endl;
+      return nullptr;
+    }
+  }
+
+  std::unique_ptr<Reader> reader;
+
+  if (suffix == "root") {
+    // Check only the first file for RNTuples
+    TFile* file = TFile::Open(filenames[0].c_str());
+    bool hasRNTuple = false;
+
+    for (auto key : *file->GetListOfKeys()) {
+      auto tkey = dynamic_cast<TKey*>(key);
+
+      // if (tkey && tkey->GetClassName() == "ROOT::Experimental::RNTuple") {
+      if (tkey && std::string(tkey->GetClassName()) == "ROOT::Experimental::RNTuple") {
+        hasRNTuple = true;
+        break;
+      }
+    }
+    if (hasRNTuple) {
+#ifdef PODIO_ENABLE_RNTUPLE
+      auto actualReader = std::make_unique<ROOTRNTupleReader>();
+      actualReader->openFiles(filenames);
+      reader = std::make_unique<Reader>(std::move(actualReader));
+#else
+      throw std::runtime_error("ROOT RNTuple reader not available. Please recompile with ROOT RNTuple support.");
+#endif
+    } else {
+      auto actualReader = std::make_unique<ROOTFrameReader>();
+      actualReader->openFiles(filenames);
+      reader = std::make_unique<Reader>(std::move(actualReader));
+    }
+  } else if (suffix == "sio") {
+#ifdef PODIO_ENABLE_SIO
+    auto actualReader = std::make_unique<SIOFrameReader>();
+    actualReader->openFiles(filenames);
+    reader = std::make_unique<Reader>(std::move(actualReader));
+#else
+    throw std::runtime_error("SIO reader not available. Please recompile with SIO support.");
+#endif
+  }
+
+  return reader;
+}
+
+} // namespace podio
diff --git a/src/Writer.cc b/src/Writer.cc
@@ -0,0 +1,52 @@
+#include "podio/Writer.h"
+
+#include "podio/ROOTFrameWriter.h"
+#ifdef PODIO_ENABLE_RNTUPLE
+  #include "podio/ROOTRNTupleWriter.h"
+#endif
+#ifdef PODIO_ENABLE_SIO
+  #include "podio/SIOFrameWriter.h"
+#endif
+
+#include <memory>
+
+namespace podio {
+
+template <typename T>
+Writer::Writer(std::unique_ptr<T> reader) : m_self(std::make_unique<WriterModel<T>>(reader.release())) {
+}
+
+std::unique_ptr<Writer> makeWriter(const std::string& filename, const std::string& type) {
+
+  auto endsWith = [](const std::string& str, const std::string& suffix) {
+    return str.size() >= suffix.size() && 0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix);
+  };
+
+  std::unique_ptr<Writer> writer;
+  if ((type == "default" && endsWith(filename, ".root")) || type == "root") {
+    std::cout << "Calling makeWriter (root)" << std::endl;
+    auto actualWriter = std::make_unique<ROOTFrameWriter>(filename);
+    writer = std::make_unique<Writer>(std::move(actualWriter));
+  } else if (type == "rntuple") {
+#ifdef PODIO_ENABLE_RNTUPLE
+    std::cout << "Calling makeWriter (rntuple)" << std::endl;
+    auto actualWriter = std::make_unique<ROOTRNTupleWriter>(filename);
+    writer = std::make_unique<Writer>(std::move(actualWriter));
+#else
+    std::cout << "ERROR: RNTuple writer not enabled" << std::endl;
+#endif
+  } else if ((type == "default" && endsWith(filename, ".sio")) || type == "sio") {
+#ifdef PODIO_ENABLE_SIO
+    std::cout << "Calling makeWriter (sio)" << std::endl;
+    auto actualWriter = std::make_unique<SIOFrameWriter>(filename);
+    writer = std::make_unique<Writer>(std::move(actualWriter));
+#else
+    std::cout << "ERROR: SIO writer not enabled" << std::endl;
+#endif
+  } else {
+    std::cout << "ERROR: Unknown writer type " << type << std::endl;
+  }
+  return writer;
+}
+
+} // namespace podio
diff --git a/tests/root_io/CMakeLists.txt b/tests/root_io/CMakeLists.txt
@@ -6,6 +6,8 @@ set(root_dependent_tests
   read_python_frame_root.cpp
   read_frame_root_multiple.cpp
   read_and_write_frame_root.cpp
+  write_interface.cpp
+  read_interface.cpp
   )
 if(ENABLE_RNTUPLE)
   set(root_dependent_tests
@@ -20,6 +22,8 @@ foreach( sourcefile ${root_dependent_tests} )
   CREATE_PODIO_TEST(${sourcefile} "${root_libs}")
 endforeach()
 
+set_property(TEST read_interface PROPERTY DEPENDS write_interface)
+
 set_tests_properties(
   read_frame_root
   read_frame_root_multiple