core: rework NumpyMap in preparation for Eigen::Tensor

include/eigenpy/numpy-map.hpp

 /*
  * Copyright 2014-2019, CNRS
- * Copyright 2018-2020, INRIA
+ * Copyright 2018-2023, INRIA
  */
 
 #ifndef __eigenpy_numpy_map_hpp__
@@ -11,42 +11,40 @@
 #include "eigenpy/stride.hpp"
 
 namespace eigenpy {
+
 template <typename MatType, typename InputScalar, int AlignmentValue,
           typename Stride, bool IsVector = MatType::IsVectorAtCompileTime>
-struct NumpyMapTraits {};
+struct numpy_map_impl_matrix;
 
-/* Wrap a numpy::array with an Eigen::Map. No memory copy. */
-template <typename MatType, typename InputScalar,
-          int AlignmentValue = EIGENPY_NO_ALIGNMENT_VALUE,
-          typename Stride = typename StrideType<MatType>::type>
-struct NumpyMap {
-  typedef NumpyMapTraits<MatType, InputScalar, AlignmentValue, Stride> Impl;
-  typedef typename Impl::EigenMap EigenMap;
+template <typename EigenType, typename InputScalar, int AlignmentValue,
+          typename Stride,
+          typename BaseType = typename get_eigen_base_type<EigenType>::type>
+struct numpy_map_impl;
 
-  static EigenMap map(PyArrayObject* pyArray, bool swap_dimensions = false);
-};
+template <typename MatType, typename InputScalar, int AlignmentValue,
+          typename Stride>
+struct numpy_map_impl<MatType, InputScalar, AlignmentValue, Stride,
+                      Eigen::MatrixBase<MatType> >
+    : numpy_map_impl_matrix<MatType, InputScalar, AlignmentValue, Stride> {};
 
-}  // namespace eigenpy
+#ifdef EIGENPY_WITH_TENSOR_SUPPORT
+template <typename TensorType, typename InputScalar, int AlignmentValue,
+          typename Stride>
+struct numpy_map_impl_tensor;
 
-/* --- DETAILS
- * ------------------------------------------------------------------ */
-/* --- DETAILS
- * ------------------------------------------------------------------ */
-/* --- DETAILS
- * ------------------------------------------------------------------ */
+#endif
 
-namespace eigenpy {
 template <typename MatType, typename InputScalar, int AlignmentValue,
           typename Stride>
-struct NumpyMapTraits<MatType, InputScalar, AlignmentValue, Stride, false> {
+struct numpy_map_impl_matrix<MatType, InputScalar, AlignmentValue, Stride,
+                             false> {
   typedef Eigen::Matrix<InputScalar, MatType::RowsAtCompileTime,
                         MatType::ColsAtCompileTime, MatType::Options>
       EquivalentInputMatrixType;
   typedef Eigen::Map<EquivalentInputMatrixType, AlignmentValue, Stride>
       EigenMap;
 
-  static EigenMap mapImpl(PyArrayObject* pyArray,
-                          bool swap_dimensions = false) {
+  static EigenMap map(PyArrayObject* pyArray, bool swap_dimensions = false) {
     enum {
       OuterStrideAtCompileTime = Stride::OuterStrideAtCompileTime,
       InnerStrideAtCompileTime = Stride::InnerStrideAtCompileTime,
@@ -135,15 +133,15 @@ struct NumpyMapTraits<MatType, InputScalar, AlignmentValue, Stride, false> {
 
 template <typename MatType, typename InputScalar, int AlignmentValue,
           typename Stride>
-struct NumpyMapTraits<MatType, InputScalar, AlignmentValue, Stride, true> {
+struct numpy_map_impl_matrix<MatType, InputScalar, AlignmentValue, Stride,
+                             true> {
   typedef Eigen::Matrix<InputScalar, MatType::RowsAtCompileTime,
                         MatType::ColsAtCompileTime, MatType::Options>
       EquivalentInputMatrixType;
   typedef Eigen::Map<EquivalentInputMatrixType, AlignmentValue, Stride>
       EigenMap;
 
-  static EigenMap mapImpl(PyArrayObject* pyArray,
-                          bool swap_dimensions = false) {
+  static EigenMap map(PyArrayObject* pyArray, bool swap_dimensions = false) {
     EIGENPY_UNUSED_VARIABLE(swap_dimensions);
     assert(PyArray_NDIM(pyArray) <= 2);
 
@@ -176,13 +174,39 @@ struct NumpyMapTraits<MatType, InputScalar, AlignmentValue, Stride, true> {
   }
 };
 
-template <typename MatType, typename InputScalar, int AlignmentValue,
+#ifdef EIGENPY_WITH_TENSOR_SUPPORT
+template <typename TensorType, typename InputScalar, int AlignmentValue,
           typename Stride>
-typename NumpyMap<MatType, InputScalar, AlignmentValue, Stride>::EigenMap
-NumpyMap<MatType, InputScalar, AlignmentValue, Stride>::map(
-    PyArrayObject* pyArray, bool swap_dimensions) {
-  return Impl::mapImpl(pyArray, swap_dimensions);
-}
+struct numpy_map_impl<TensorType, InputScalar, AlignmentValue, Stride,
+                      Eigen::TensorBase<TensorType> > {
+  typedef TensorType Tensor;
+  typedef typename Eigen::internal::traits<TensorType>::Index Index;
+  static const Index NumIndices = TensorType::NumIndices;
+
+  typedef Eigen::Tensor<InputScalar, NumIndices, Tensor::Options, Index>
+      EquivalentInputTensorType;
+  typedef typename EquivalentInputTensorType::Dimensions Dimensions;
+  typedef Eigen::TensorMap<EquivalentInputTensorType, Tensor::Options> EigenMap;
+
+  static EigenMap map(PyArrayObject* pyArray, bool swap_dimensions = false) {
+    assert(PyArray_NDIM(pyArray) == NumIndices || NumIndices == Eigen::Dynamic);
+
+    Eigen::DSizes<Index, NumIndices> dimensions(PyArray_NDIM(pyArray));
+    for (int k = 0; k < PyArray_NDIM(pyArray); ++k)
+      dimensions[k] = PyArray_DIMS(pyArray)[k];
+
+    InputScalar* pyData = reinterpret_cast<InputScalar*>(PyArray_DATA(pyArray));
+    return EigenMap(pyData, dimensions);
+  }
+};
+#endif
+
+/* Wrap a numpy::array with an Eigen::Map. No memory copy. */
+template <typename EigenType, typename InputScalar,
+          int AlignmentValue = EIGENPY_NO_ALIGNMENT_VALUE,
+          typename Stride = typename StrideType<EigenType>::type>
+struct NumpyMap
+    : numpy_map_impl<EigenType, InputScalar, AlignmentValue, Stride> {};
 
 }  // namespace eigenpy