void DefineCircularBuffer()

in rlmeta/cc/circular_buffer.cc [88:126]

• 39 lines of code
• 1 McCabe index (conditional complexity)

void DefineCircularBuffer(py::module& m) {
  py::class_<CircularBuffer, std::shared_ptr<CircularBuffer>>(m,
                                                              "CircularBuffer")
      .def(py::init<int64_t>())
      .def("__len__", &CircularBuffer::Size)
      .def("size", &CircularBuffer::Size)
      .def("capacity", &CircularBuffer::capacity)
      .def("cursor", &CircularBuffer::cursor)
      .def("__getitem__",
           py::overload_cast<int64_t>(&CircularBuffer::At, py::const_))
      .def("__getitem__", py::overload_cast<const py::array_t<int64_t>&>(
                              &CircularBuffer::At, py::const_))
      .def("__getitem__", py::overload_cast<const torch::Tensor&>(
                              &CircularBuffer::At, py::const_))
      .def("at", py::overload_cast<int64_t>(&CircularBuffer::At, py::const_))
      .def("at", py::overload_cast<const py::array_t<int64_t>&>(
                     &CircularBuffer::At, py::const_))
      .def("at", py::overload_cast<const torch::Tensor&>(&CircularBuffer::At,
                                                         py::const_))
      .def("reset", &CircularBuffer::Reset)
      .def("append", &CircularBuffer::Append)
      .def("extend",
           py::overload_cast<const py::tuple&>(&CircularBuffer::Extend))
      .def("extend",
           py::overload_cast<const py::list&>(&CircularBuffer::Extend))
      .def(py::pickle(
          [](const CircularBuffer& s) {
            return py::make_tuple(s.capacity(), s.DumpData(), s.cursor());
          },
          [](const py::tuple& t) {
            assert(t.size() == 3);
            const int64_t capacity = t[0].cast<int64_t>();
            const py::tuple data = py::reinterpret_borrow<py::tuple>(t[1]);
            const int64_t cursor = t[2].cast<int64_t>();
            CircularBuffer buffer(capacity);
            buffer.LoadData(data, cursor);
            return buffer;
          }));
}