buffer_info.h 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. /*
  2. pybind11/buffer_info.h: Python buffer object interface
  3. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  4. All rights reserved. Use of this source code is governed by a
  5. BSD-style license that can be found in the LICENSE file.
  6. */
  7. #pragma once
  8. #include "detail/common.h"
  9. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  10. PYBIND11_NAMESPACE_BEGIN(detail)
  11. // Default, C-style strides
  12. inline std::vector<ssize_t> c_strides(const std::vector<ssize_t> &shape, ssize_t itemsize) {
  13. auto ndim = shape.size();
  14. std::vector<ssize_t> strides(ndim, itemsize);
  15. if (ndim > 0) {
  16. for (size_t i = ndim - 1; i > 0; --i) {
  17. strides[i - 1] = strides[i] * shape[i];
  18. }
  19. }
  20. return strides;
  21. }
  22. // F-style strides; default when constructing an array_t with `ExtraFlags & f_style`
  23. inline std::vector<ssize_t> f_strides(const std::vector<ssize_t> &shape, ssize_t itemsize) {
  24. auto ndim = shape.size();
  25. std::vector<ssize_t> strides(ndim, itemsize);
  26. for (size_t i = 1; i < ndim; ++i) {
  27. strides[i] = strides[i - 1] * shape[i - 1];
  28. }
  29. return strides;
  30. }
  31. template <typename T, typename SFINAE = void>
  32. struct compare_buffer_info;
  33. PYBIND11_NAMESPACE_END(detail)
  34. /// Information record describing a Python buffer object
  35. struct buffer_info {
  36. void *ptr = nullptr; // Pointer to the underlying storage
  37. ssize_t itemsize = 0; // Size of individual items in bytes
  38. ssize_t size = 0; // Total number of entries
  39. std::string format; // For homogeneous buffers, this should be set to
  40. // format_descriptor<T>::format()
  41. ssize_t ndim = 0; // Number of dimensions
  42. std::vector<ssize_t> shape; // Shape of the tensor (1 entry per dimension)
  43. std::vector<ssize_t> strides; // Number of bytes between adjacent entries
  44. // (for each per dimension)
  45. bool readonly = false; // flag to indicate if the underlying storage may be written to
  46. buffer_info() = default;
  47. buffer_info(void *ptr,
  48. ssize_t itemsize,
  49. const std::string &format,
  50. ssize_t ndim,
  51. detail::any_container<ssize_t> shape_in,
  52. detail::any_container<ssize_t> strides_in,
  53. bool readonly = false)
  54. : ptr(ptr), itemsize(itemsize), size(1), format(format), ndim(ndim),
  55. shape(std::move(shape_in)), strides(std::move(strides_in)), readonly(readonly) {
  56. if (ndim != (ssize_t) shape.size() || ndim != (ssize_t) strides.size()) {
  57. pybind11_fail("buffer_info: ndim doesn't match shape and/or strides length");
  58. }
  59. for (size_t i = 0; i < (size_t) ndim; ++i) {
  60. size *= shape[i];
  61. }
  62. }
  63. template <typename T>
  64. buffer_info(T *ptr,
  65. detail::any_container<ssize_t> shape_in,
  66. detail::any_container<ssize_t> strides_in,
  67. bool readonly = false)
  68. : buffer_info(private_ctr_tag(),
  69. ptr,
  70. sizeof(T),
  71. format_descriptor<T>::format(),
  72. static_cast<ssize_t>(shape_in->size()),
  73. std::move(shape_in),
  74. std::move(strides_in),
  75. readonly) {}
  76. buffer_info(void *ptr,
  77. ssize_t itemsize,
  78. const std::string &format,
  79. ssize_t size,
  80. bool readonly = false)
  81. : buffer_info(ptr, itemsize, format, 1, {size}, {itemsize}, readonly) {}
  82. template <typename T>
  83. buffer_info(T *ptr, ssize_t size, bool readonly = false)
  84. : buffer_info(ptr, sizeof(T), format_descriptor<T>::format(), size, readonly) {}
  85. template <typename T>
  86. buffer_info(const T *ptr, ssize_t size, bool readonly = true)
  87. : buffer_info(
  88. const_cast<T *>(ptr), sizeof(T), format_descriptor<T>::format(), size, readonly) {}
  89. explicit buffer_info(Py_buffer *view, bool ownview = true)
  90. : buffer_info(
  91. view->buf,
  92. view->itemsize,
  93. view->format,
  94. view->ndim,
  95. {view->shape, view->shape + view->ndim},
  96. /* Though buffer::request() requests PyBUF_STRIDES, ctypes objects
  97. * ignore this flag and return a view with NULL strides.
  98. * When strides are NULL, build them manually. */
  99. view->strides
  100. ? std::vector<ssize_t>(view->strides, view->strides + view->ndim)
  101. : detail::c_strides({view->shape, view->shape + view->ndim}, view->itemsize),
  102. (view->readonly != 0)) {
  103. // NOLINTNEXTLINE(cppcoreguidelines-prefer-member-initializer)
  104. this->m_view = view;
  105. // NOLINTNEXTLINE(cppcoreguidelines-prefer-member-initializer)
  106. this->ownview = ownview;
  107. }
  108. buffer_info(const buffer_info &) = delete;
  109. buffer_info &operator=(const buffer_info &) = delete;
  110. buffer_info(buffer_info &&other) noexcept { (*this) = std::move(other); }
  111. buffer_info &operator=(buffer_info &&rhs) noexcept {
  112. ptr = rhs.ptr;
  113. itemsize = rhs.itemsize;
  114. size = rhs.size;
  115. format = std::move(rhs.format);
  116. ndim = rhs.ndim;
  117. shape = std::move(rhs.shape);
  118. strides = std::move(rhs.strides);
  119. std::swap(m_view, rhs.m_view);
  120. std::swap(ownview, rhs.ownview);
  121. readonly = rhs.readonly;
  122. return *this;
  123. }
  124. ~buffer_info() {
  125. if (m_view && ownview) {
  126. PyBuffer_Release(m_view);
  127. delete m_view;
  128. }
  129. }
  130. Py_buffer *view() const { return m_view; }
  131. Py_buffer *&view() { return m_view; }
  132. /* True if the buffer item type is equivalent to `T`. */
  133. // To define "equivalent" by example:
  134. // `buffer_info::item_type_is_equivalent_to<int>(b)` and
  135. // `buffer_info::item_type_is_equivalent_to<long>(b)` may both be true
  136. // on some platforms, but `int` and `unsigned` will never be equivalent.
  137. // For the ground truth, please inspect `detail::compare_buffer_info<>`.
  138. template <typename T>
  139. bool item_type_is_equivalent_to() const {
  140. return detail::compare_buffer_info<T>::compare(*this);
  141. }
  142. private:
  143. struct private_ctr_tag {};
  144. buffer_info(private_ctr_tag,
  145. void *ptr,
  146. ssize_t itemsize,
  147. const std::string &format,
  148. ssize_t ndim,
  149. detail::any_container<ssize_t> &&shape_in,
  150. detail::any_container<ssize_t> &&strides_in,
  151. bool readonly)
  152. : buffer_info(
  153. ptr, itemsize, format, ndim, std::move(shape_in), std::move(strides_in), readonly) {}
  154. Py_buffer *m_view = nullptr;
  155. bool ownview = false;
  156. };
  157. PYBIND11_NAMESPACE_BEGIN(detail)
  158. template <typename T, typename SFINAE>
  159. struct compare_buffer_info {
  160. static bool compare(const buffer_info &b) {
  161. // NOLINTNEXTLINE(bugprone-sizeof-expression) Needed for `PyObject *`
  162. return b.format == format_descriptor<T>::format() && b.itemsize == (ssize_t) sizeof(T);
  163. }
  164. };
  165. template <typename T>
  166. struct compare_buffer_info<T, detail::enable_if_t<std::is_integral<T>::value>> {
  167. static bool compare(const buffer_info &b) {
  168. return (size_t) b.itemsize == sizeof(T)
  169. && (b.format == format_descriptor<T>::value
  170. || ((sizeof(T) == sizeof(long))
  171. && b.format == (std::is_unsigned<T>::value ? "L" : "l"))
  172. || ((sizeof(T) == sizeof(size_t))
  173. && b.format == (std::is_unsigned<T>::value ? "N" : "n")));
  174. }
  175. };
  176. PYBIND11_NAMESPACE_END(detail)
  177. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)