Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion include/xtensor-python/pyarray.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,20 @@ namespace pybind11
{
using type = xt::pyarray<T>;

bool load(handle src, bool)
bool load(handle src, bool convert)
{
if (!convert)
{
if (!PyArray_Check(src.ptr()))
{
return false;
}
int type_num = xt::detail::numpy_traits<T>::type_num;
if (PyArray_TYPE(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
{
return false;
}
}
value = type::ensure(src);
return static_cast<bool>(value);
}
Expand Down
15 changes: 14 additions & 1 deletion include/xtensor-python/pytensor.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,21 @@ namespace pybind11
{
using type = xt::pytensor<T, N>;

bool load(handle src, bool)
bool load(handle src, bool convert)
{
if (!convert)
{
if (!PyArray_Check(src.ptr()))
{
return false;
}
int type_num = xt::detail::numpy_traits<T>::type_num;
if (PyArray_TYPE(reinterpret_cast<PyArrayObject*>(src.ptr())) != type_num)
{
return false;
}
}

value = type::ensure(src);
return static_cast<bool>(value);
}
Expand Down
24 changes: 24 additions & 0 deletions test_python/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,25 @@ double readme_example2(double i, double j)
return std::sin(i) - std::cos(j);
}

auto complex_overload(const xt::pyarray<std::complex<double>>& a)
{
return a;
}
auto no_complex_overload(const xt::pyarray<double>& a)
{
return a;
}

auto complex_overload_reg(const std::complex<double>& a)
{
return a;
}

auto no_complex_overload_reg(const double& a)
{
return a;
}

// Vectorize Examples

int add(int i, int j)
Expand All @@ -58,6 +77,11 @@ PYBIND11_PLUGIN(xtensor_python_test)
m.def("example1", example1);
m.def("example2", example2);

m.def("complex_overload", no_complex_overload);
m.def("complex_overload", complex_overload);
m.def("complex_overload_reg", no_complex_overload_reg);
m.def("complex_overload_reg", complex_overload_reg);

m.def("readme_example1", readme_example1);
m.def("readme_example2", xt::pyvectorize(readme_example2));

Expand Down
17 changes: 17 additions & 0 deletions test_python/test_pyarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,23 @@ def test_readme_example1(self):
y = xt.readme_example1(v)
np.testing.assert_allclose(y, 1.2853996391883833, 1e-12)

def test_complex_overload_reg(self):
a = 23.23
c = 2.0 + 3.1j
self.assertEqual(xt.complex_overload_reg(a), a)
self.assertEqual(xt.complex_overload_reg(c), c)

def test_complex_overload(self):
a = np.random.rand(3, 3)
b = np.random.rand(3, 3)
c = a + b * 1j
y = xt.complex_overload(c)
np.testing.assert_allclose(np.imag(y), np.imag(c))
np.testing.assert_allclose(np.real(y), np.real(c))
x = xt.complex_overload(b)
self.assertEqual(x.dtype, b.dtype)
np.testing.assert_allclose(x, b)

def test_readme_example2(self):
x = np.arange(15).reshape(3, 5)
y = [1, 2, 3, 4, 5]
Expand Down