include/pybind11/complex.h

0001 /*
0002     pybind11/complex.h: Complex number support
0003
0004     Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
0005
0006     All rights reserved. Use of this source code is governed by a
0007     BSD-style license that can be found in the LICENSE file.
0008 */
0009
0010 #pragma once
0011
0012 #include "pybind11.h"
0013
0014 #include <complex>
0015
0016 /// glibc defines I as a macro which breaks things, e.g., boost template names
0017 #ifdef I
0018 #    undef I
0019 #endif
0020
0021 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
0022
0023 template <typename T>
0024 struct format_descriptor<std::complex<T>, detail::enable_if_t<std::is_floating_point<T>::value>> {
0025     static constexpr const char c = format_descriptor<T>::c;
0026     static constexpr const char value[3] = {'Z', c, '\0'};
0027     static std::string format() { return std::string(value); }
0028 };
0029
0030 #ifndef PYBIND11_CPP17
0031
0032 template <typename T>
0033 constexpr const char
0034     format_descriptor<std::complex<T>,
0035                       detail::enable_if_t<std::is_floating_point<T>::value>>::value[3];
0036
0037 #endif
0038
0039 PYBIND11_NAMESPACE_BEGIN(detail)
0040
0041 template <typename T>
0042 struct is_fmt_numeric<std::complex<T>, detail::enable_if_t<std::is_floating_point<T>::value>> {
0043     static constexpr bool value = true;
0044     static constexpr int index = is_fmt_numeric<T>::index + 3;
0045 };
0046
0047 template <typename T>
0048 class type_caster<std::complex<T>> {
0049 public:
0050     bool load(handle src, bool convert) {
0051         if (!src) {
0052             return false;
0053         }
0054         if (!convert && !PyComplex_Check(src.ptr())) {
0055             return false;
0056         }
0057         handle src_or_index = src;
0058         // PyPy: 7.3.7's 3.8 does not implement PyLong_*'s __index__ calls.
0059         // The same logic is used in numeric_caster for ints and floats
0060 #if defined(PYPY_VERSION)
0061         object index;
0062         if (PYBIND11_INDEX_CHECK(src.ptr())) {
0063             index = reinterpret_steal<object>(PyNumber_Index(src.ptr()));
0064             if (!index) {
0065                 PyErr_Clear();
0066                 if (!convert)
0067                     return false;
0068             } else {
0069                 src_or_index = index;
0070             }
0071         }
0072 #endif
0073         Py_complex result = PyComplex_AsCComplex(src_or_index.ptr());
0074         if (result.real == -1.0 && PyErr_Occurred()) {
0075             PyErr_Clear();
0076             return false;
0077         }
0078         value = std::complex<T>((T) result.real, (T) result.imag);
0079         return true;
0080     }
0081
0082     static handle
0083     cast(const std::complex<T> &src, return_value_policy /* policy */, handle /* parent */) {
0084         return PyComplex_FromDoubles((double) src.real(), (double) src.imag());
0085     }
0086
0087     PYBIND11_TYPE_CASTER(
0088         std::complex<T>,
0089         io_name("typing.SupportsComplex | typing.SupportsFloat | typing.SupportsIndex",
0090                 "complex"));
0091 };
0092 PYBIND11_NAMESPACE_END(detail)
0093 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)