Skip to content

Commit afb4a2a

Browse files
committed
Merge remote-tracking branch 'upstream/master' into fix-py313t
2 parents 9a039df + 5f2c678 commit afb4a2a

File tree

4 files changed

+78
-0
lines changed

4 files changed

+78
-0
lines changed

include/pybind11/detail/common.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,10 @@
252252
# define PYBIND11_HAS_U8STRING 1
253253
#endif
254254

255+
#if defined(PYBIND11_CPP20) && defined(__cpp_lib_span) && __cpp_lib_span >= 202002L
256+
# define PYBIND11_HAS_SPAN 1
257+
#endif
258+
255259
// See description of PR #4246:
256260
#if !defined(PYBIND11_NO_ASSERT_GIL_HELD_INCREF_DECREF) && !defined(NDEBUG) \
257261
&& !defined(PYPY_VERSION) && !defined(PYBIND11_ASSERT_GIL_HELD_INCREF_DECREF)

include/pybind11/numpy.h

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@
2929
#include <utility>
3030
#include <vector>
3131

32+
#ifdef PYBIND11_HAS_SPAN
33+
# include <span>
34+
#endif
35+
3236
#if defined(PYBIND11_NUMPY_1_ONLY)
3337
# error "PYBIND11_NUMPY_1_ONLY is no longer supported (see PR #5595)."
3438
#endif
@@ -1143,6 +1147,13 @@ class array : public buffer {
11431147
/// Dimensions of the array
11441148
const ssize_t *shape() const { return detail::array_proxy(m_ptr)->dimensions; }
11451149

1150+
#ifdef PYBIND11_HAS_SPAN
1151+
/// Dimensions of the array as a span
1152+
std::span<const ssize_t, std::dynamic_extent> shape_span() const {
1153+
return std::span(shape(), static_cast<std::size_t>(ndim()));
1154+
}
1155+
#endif
1156+
11461157
/// Dimension along a given axis
11471158
ssize_t shape(ssize_t dim) const {
11481159
if (dim >= ndim()) {
@@ -1154,6 +1165,13 @@ class array : public buffer {
11541165
/// Strides of the array
11551166
const ssize_t *strides() const { return detail::array_proxy(m_ptr)->strides; }
11561167

1168+
#ifdef PYBIND11_HAS_SPAN
1169+
/// Strides of the array as a span
1170+
std::span<const ssize_t, std::dynamic_extent> strides_span() const {
1171+
return std::span(strides(), static_cast<std::size_t>(ndim()));
1172+
}
1173+
#endif
1174+
11571175
/// Stride along a given axis
11581176
ssize_t strides(ssize_t dim) const {
11591177
if (dim >= ndim()) {

tests/test_numpy_array.cpp

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
#include <cstdint>
1616
#include <utility>
17+
#include <vector>
1718

1819
// Size / dtype checks.
1920
struct DtypeCheck {
@@ -246,6 +247,22 @@ TEST_SUBMODULE(numpy_array, sm) {
246247
sm.def("nbytes", [](const arr &a) { return a.nbytes(); });
247248
sm.def("owndata", [](const arr &a) { return a.owndata(); });
248249

250+
#ifdef PYBIND11_HAS_SPAN
251+
// test_shape_strides_span
252+
sm.def("shape_span", [](const arr &a) {
253+
auto span = a.shape_span();
254+
return std::vector<py::ssize_t>(span.begin(), span.end());
255+
});
256+
sm.def("strides_span", [](const arr &a) {
257+
auto span = a.strides_span();
258+
return std::vector<py::ssize_t>(span.begin(), span.end());
259+
});
260+
// Test that spans can be used to construct new arrays
261+
sm.def("array_from_spans", [](const arr &a) {
262+
return py::array(a.dtype(), a.shape_span(), a.strides_span(), a.data(), a);
263+
});
264+
#endif
265+
249266
// test_index_offset
250267
def_index_fn(index_at, const arr &);
251268
def_index_fn(index_at_t, const arr_t &);

tests/test_numpy_array.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,45 @@ def test_array_attributes():
6868
assert not m.owndata(a)
6969

7070

71+
@pytest.mark.skipif(not hasattr(m, "shape_span"), reason="std::span not available")
72+
def test_shape_strides_span():
73+
# Test 0-dimensional array (scalar)
74+
a = np.array(42, "f8")
75+
assert m.ndim(a) == 0
76+
assert m.shape_span(a) == []
77+
assert m.strides_span(a) == []
78+
79+
# Test 1-dimensional array
80+
a = np.array([1, 2, 3, 4], "u2")
81+
assert m.ndim(a) == 1
82+
assert m.shape_span(a) == [4]
83+
assert m.strides_span(a) == [2]
84+
85+
# Test 2-dimensional array
86+
a = np.array([[1, 2, 3], [4, 5, 6]], "u2").view()
87+
a.flags.writeable = False
88+
assert m.ndim(a) == 2
89+
assert m.shape_span(a) == [2, 3]
90+
assert m.strides_span(a) == [6, 2]
91+
92+
# Test 3-dimensional array
93+
a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], "i4")
94+
assert m.ndim(a) == 3
95+
assert m.shape_span(a) == [2, 2, 2]
96+
# Verify spans match regular shape/strides
97+
assert list(m.shape_span(a)) == list(m.shape(a))
98+
assert list(m.strides_span(a)) == list(m.strides(a))
99+
100+
# Test that spans can be used to construct new arrays
101+
original = np.array([[1, 2, 3], [4, 5, 6]], "f4")
102+
new_array = m.array_from_spans(original)
103+
assert new_array.shape == original.shape
104+
assert new_array.strides == original.strides
105+
assert new_array.dtype == original.dtype
106+
# Verify data is shared (since we pass the same data pointer)
107+
np.testing.assert_array_equal(new_array, original)
108+
109+
71110
@pytest.mark.parametrize(
72111
("args", "ret"), [([], 0), ([0], 0), ([1], 3), ([0, 1], 1), ([1, 2], 5)]
73112
)

0 commit comments

Comments
 (0)