Repository: apytypes/apytypes (language: python)

Issue:
C++: hoist common member functions out of arrays
This is an incomplete list of candidate member functions in `APyFixedArray` and `APyFloatArray` that are near identical source-code-wise, and should probably be hoisted to a common array base class.

* [x] `transpose`
* [x] `rehsape`
* [x] `flatten`
* [x] `ravel`
* [x] `swapaxes` 
* [x] `squeeze`
* [x] `__getitem__` (+ family of helpers)
* [x] `size`
* [x] `ndim`
* [x] `python_get_shape`
* [x] `broadcast_to`
* [x] `broadcast_to_python`
* [ ] ~~`identity`~~
* [ ] ~~`eye`~~
* [x] `diagonal`
* [x] `full`
* [ ] ~~`zeros`~~
* [ ] ~~`ones`~~



Discussion / hints:
* #519 
  * `size`
  * `ndim`
  * `python_get_shape`
* #527 
  * `__getitem__` (+ family of helpers)
* #529 
  * `broadcast_to`
  * `broadcast_to_python`
* #531 
  * `reshape`

Fix patch:
diff --git a/src/apyarray.h b/src/apyarray.h
index 19793eceb..a358527a4 100644
--- a/src/apyarray.h
+++ b/src/apyarray.h
@@ -1,5 +1,6 @@
 /*
- * APyTypes array base class, for sharing functionality between array types.
+ * APyTypes array base class, for sharing functionality between array types. All the
+ * member functions of the APyTypes array types I don't want to write twice...
  */
 
 #ifndef _APYARRAY_H
@@ -7,6 +8,7 @@
 
 #include "apybuffer.h"
 #include "apytypes_util.h"
+#include "array_utils.h"
 #include "broadcast.h"
 
 // Python object access through Nanobind
@@ -15,6 +17,7 @@
 namespace nb = nanobind;
 
 #include <iterator> // std::begin
+#include <set>      // std::set
 #include <variant>  // std::variant
 
 template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T> {
@@ -301,7 +304,7 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
         }
     }
 
-    //! Top-level Python-exported `__getitem__` function
+    //! Top-level Python exported `__getitem__` function
     auto get_item(std::variant<nb::int_, nb::slice, nb::ellipsis, nb::tuple> key) const
     {
         using SCALAR_TYPE = typename ARRAY_TYPE::SCALAR_VARIANT;
@@ -440,6 +443,130 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
         return result;
     }
 
-}; // class APyArray
+    //! Python exported `flatten` method
+    ARRAY_TYPE flatten() const
+    {
+        // Reuse the reshape function to flatten the array
+        return reshape(nb::make_tuple(-1));
+    }
+
+    //! Python exported `ravel` method
+    ARRAY_TYPE ravel() const { return flatten(); }
+
+    //! Python exported `transpose` method.
+    ARRAY_TYPE transpose(std::optional<nb::tuple> axes = std::nullopt) const
+    {
+        switch (_ndim) {
+        case 0:
+        case 1:
+            // Behave like NumPy, simply return `*this` if single dimensional
+            return *static_cast<const ARRAY_TYPE*>(this);
+        case 2: {
+            // Optimized code for `_ndim` == 2
+            ARRAY_TYPE result
+                = static_cast<const ARRAY_TYPE*>(this)->create_array(_shape);
+            std::reverse(result._shape.begin(), result._shape.end());
+            for (std::size_t y = 0; y < _shape[0]; y++) {
+                for (std::size_t x = 0; x < _shape[1]; x++) {
+                    std::copy_n(
+                        _data.begin() + (y * _shape[1] + x) * _itemsize,       // src
+                        _itemsize,                                             // limbs
+                        result._data.begin() + (x * _shape[0] + y) * _itemsize // dst
+                    );
+                }
+            }
+            return result;
+        }
+
+        default: {
+            std::vector<size_t> new_axis(_ndim);
+
+            if (axes.has_value()) {
+                std::variant<nb::tuple, nb::int_> axis_variant = axes.value();
+                new_axis = get_normalized_axes(axis_variant, _ndim);
+            } else {
+                // reverse the order of axes by default
+                std::iota(new_axis.begin(), new_axis.end(), 0);
+                std::reverse(new_axis.begin(), new_axis.end());
+            }
+
+            std::vector<size_t> new_shape(_ndim);
+            for (size_t i = 0; i < _ndim; ++i) {
+                new_shape[i] = _shape[new_axis[i]];
+            }
+
+            ARRAY_TYPE result
+                = static_cast<const ARRAY_TYPE*>(this)->create_array(new_shape);
+            transpose_axes_and_copy_data(
+                _data.begin(), result._data.begin(), _shape, new_axis, _itemsize
+            );
+            return result;
+        }
+        }
+    }
+
+    //! Python exported `squeeze` method
+    ARRAY_TYPE squeeze(std::optional<std::variant<nb::int_, nb::tuple>> axis) const
+    {
+        std::vector<std::size_t> shape = _shape;
+        std::set<std::ptrdiff_t> axis_set;
+        if (axis.has_value()) {
+            // Given an int or tuple of ints, remove the specified dimensions if their
+            // size is one. Else throw an error.
+            auto ax = *axis;
+            nb::tuple axis_tuple;
+
+            if (std::holds_alternative<nb::tuple>(ax)) {
+                axis_tuple = std::get<nb::tuple>(ax);
+            } else if (std::holds_alternative<nb::int_>(ax)) {
+                std::ptrdiff_t axis = std::ptrdiff_t(std::get<nb::int_>(ax));
+                axis_tuple = nb::make_tuple(axis);
+            }
+            for (auto ptr = axis_tuple.begin(); ptr != axis_tuple.end(); ptr++) {
+                std::ptrdiff_t axis_n = std::ptrdiff_t(nanobind::cast<nb::int_>(*ptr));
+                if (axis_n >= std::ptrdiff_t(shape.size())) {
+                    std::string error_msg = fmt::format(
+                        "{}.squeeze: specified axis larger than number of array "
+                        "dimensions",
+                        ARRAY_TYPE::ARRAY_NAME
+                    );
+                    throw nb::index_error(error_msg.c_str());
+                }
+                axis_set.insert(axis_n);
+            }
+            std::ptrdiff_t cnt = 0;
+            auto predicate = [&](std::size_t dim) {
+                if (axis_set.find(cnt) != axis_set.end() && dim != 1) {
+                    std::string error_msg = fmt::format(
+                        "{}.squeeze: cannot squeeze non-one dimensions",
+                        ARRAY_TYPE::ARRAY_NAME
+                    );
+                    throw nb::value_error(error_msg.c_str());
+                }
+                return axis_set.find(cnt++) != axis_set.end() && dim == 1;
+            };
+            shape.erase(
+                std::remove_if(shape.begin(), shape.end(), predicate), shape.end()
+            );
+        } else {
+            // Given no specified axis, remove all dimensions of size one.
+            auto predicate = [](std::size_t dim) { return dim == 1; };
+            shape.erase(
+                std::remove_if(shape.begin(), shape.end(), predicate), shape.end()
+            );
+        }
+
+        // Don't squeeze all dimensions...
+        if (shape.size() == 0) {
+            shape = { 1 };
+        }
+
+        // Create resulting array and copy the data
+        ARRAY_TYPE result = static_cast<const ARRAY_TYPE*>(this)->create_array(shape);
+        std::copy_n(_data.begin(), _data.size(), result._data.begin());
+        return result;
+    }
+
+}; // end class: `APyArray`
 
 #endif
diff --git a/src/apyfixedarray.cc b/src/apyfixedarray.cc
index 02c96c4ba..c83d7d58c 100644
--- a/src/apyfixedarray.cc
+++ b/src/apyfixedarray.cc
@@ -757,114 +757,10 @@ APyFixedArray APyFixedArray::swapaxes(nb::int_ axis1, nb::int_ axis2) const
     return result;
 }
 
-APyFixedArray APyFixedArray::transpose(std::optional<nb::tuple> axes) const
-{
-    switch (_ndim) {
-    case 0:
-    case 1:
-        // Behave like `NumPy`, simply returns `*this` if single-dimensional
-        return *this;
-    case 2: {
-        // Optimized code for dim == 2
-        // Resulting array with shape dimensions
-        APyFixedArray result(_shape, bits(), int_bits());
-        std::reverse(result._shape.begin(), result._shape.end());
-        // Copy data
-        for (std::size_t y = 0; y < _shape[0]; y++) {
-            for (std::size_t x = 0; x < _shape[1]; x++) {
-                std::copy_n(
-                    _data.begin() + (y * _shape[1] + x) * _itemsize,       // src
-                    _itemsize,                                             // limbs
-                    result._data.begin() + (x * _shape[0] + y) * _itemsize // dst
-                );
-            }
-        }
-        return result;
-    }
-
-    default: {
-        std::vector<size_t> new_axis(_ndim);
-
-        if (axes.has_value()) {
-            std::variant<nb::tuple, nb::int_> axis_variant = axes.value();
-            new_axis = get_normalized_axes(axis_variant, _ndim);
-        } else {
-            // reverse the order of axes by default
-            std::iota(new_axis.begin(), new_axis.end(), 0);
-            std::reverse(new_axis.begin(), new_axis.end());
-        }
-
-        std::vector<size_t> new_shape(_ndim);
-        for (size_t i = 0; i < _ndim; ++i) {
-            new_shape[i] = _shape[new_axis[i]];
-        }
-
-        APyFixedArray result(new_shape, bits(), int_bits());
-        ::transpose_axes_and_copy_data(
-            _data.begin(), result._data.begin(), _shape, new_axis, _itemsize
-        );
-        return result;
-    }
-    }
-}
-
 /* ********************************************************************************** *
  * *                               Other methods                                    * *
  * ********************************************************************************** */
 
-APyFixedArray
-APyFixedArray::squeeze(std::optional<std::variant<nb::int_, nb::tuple>> axis) const
-{
-    std::vector<std::size_t> shape = _shape;
-    std::set<int> axis_set;
-    if (axis.has_value()) {
-        // given an int or tuple of ints, remove the specified dimensions if size is one
-        // or throw an error
-        auto ax = axis.value();
-        nb::tuple axis_tuple;
-
-        if (std::holds_alternative<nb::tuple>(ax)) {
-            axis_tuple = std::get<nb::tuple>(ax);
-        } else if (std::holds_alternative<nb::int_>(ax)) {
-            int axis = int(std::get<nb::int_>(ax));
-            axis_tuple = nb::make_tuple(axis);
-        }
-        for (auto ptr = axis_tuple.begin(); ptr != axis_tuple.end(); ptr++) {
-            int axis_n = int(nanobind::cast<nb::int_>(*ptr));
-            if (axis_n >= int(shape.size())) {
-                throw nb::index_error(
-                    "Specified axis with larger than number of dimensions in the "
-                    "APyFixedArray"
-                );
-            }
-            axis_set.insert(axis_n);
-        }
-        int cnt = 0;
-        auto predicate = [&](std::size_t dim) {
-            if (axis_set.find(cnt) != axis_set.end() && dim != 1) {
-                throw nb::value_error(
-                    "Cannot select an axis to squeeze that has size other than one"
-                );
-            }
-            return axis_set.find(cnt++) != axis_set.end() && dim == 1;
-        };
-        shape.erase(std::remove_if(shape.begin(), shape.end(), predicate), shape.end());
-    } else {
-        // Given no specified axis, remove all dimensions of size one.
-        auto predicate = [](std::size_t dim) { return dim == 1; };
-        shape.erase(std::remove_if(shape.begin(), shape.end(), predicate), shape.end());
-    }
-    // if all dimensions where one, set the dimension to one, to avoid an error
-    if (shape.size() == 0) {
-        shape = { 1 };
-    }
-
-    // create the resulting Array and copy the data
-    APyFixedArray result(shape, bits(), int_bits());
-    std::copy_n(_data.begin(), _data.size(), result._data.begin());
-    return result;
-}
-
 APyFixedArray APyFixedArray::cumulative_prod_sum_function(
     void (*pos_func)(
         std::size_t,
@@ -1489,18 +1385,6 @@ std::string APyFixedArray::repr() const
     return ss.str();
 }
 
-APyFixedArray APyFixedArray::flatten() const
-{
-    // Reuse the reshape function to flatten the array
-    return this->reshape(nb::make_tuple(-1));
-}
-
-APyFixedArray APyFixedArray::ravel() const
-{
-    // same as flatten as for now
-    return this->flatten();
-}
-
 APyFixedArray APyFixedArray::abs() const
 {
     // Increase word length of result by one
diff --git a/src/apyfixedarray.h b/src/apyfixedarray.h
index 8583ec3ce..e2f2ee397 100644
--- a/src/apyfixedarray.h
+++ b/src/apyfixedarray.h
@@ -186,21 +186,10 @@ class APyFixedArray : public APyArray<mp_limb_t, APyFixedArray> {
     //! Swaps the positions of two axes in the array.
     APyFixedArray swapaxes(nb::int_ axis1, nb::int_ axis2) const;
 
-    //! Transposition function. For a 1-D array, Return an exact copy of `*this`. For
-    //! a 2-D array, Return the matrix transposition of `*this`.
-    APyFixedArray transpose(std::optional<nb::tuple> axes = std::nullopt) const;
-
     /* ****************************************************************************** *
      * *                          Public member functions                           * *
      * ****************************************************************************** */
 
-    //! Broadcast to a new shape
-    // APyFixedArray broadcast_to(const std::vector<std::size_t> shape) const;
-
-    //! Squeeze this array. Return new `APyFixedArray`.
-    APyFixedArray
-    squeeze(std::optional<std::variant<nb::int_, nb::tuple>> axis = std::nullopt) const;
-
     //! Perform a linear convolution with `other` using `mode`
     APyFixedArray convolve(const APyFixedArray& other, const std::string& mode) const;
 
@@ -251,12 +240,6 @@ class APyFixedArray : public APyArray<mp_limb_t, APyFixedArray> {
     //! Python `__repr__()` function
     std::string repr() const;
 
-    //! Reshape flatten into 1d shape
-    APyFixedArray flatten() const;
-
-    //! Same as flatten as for now
-    APyFixedArray ravel() const;
-
     //! Number of bits
     APY_INLINE int bits() const noexcept { return _bits; }
 
diff --git a/src/apyfloatarray.cc b/src/apyfloatarray.cc
index 1f88f1d0f..ef35784c7 100644
--- a/src/apyfloatarray.cc
+++ b/src/apyfloatarray.cc
@@ -1109,62 +1109,6 @@ APyFloatArray APyFloatArray::arange(
     return result;
 }
 
-/* ****************************************************************************** *
- *                            Methods for array modification                      *
- * ****************************************************************************** */
-
-APyFloatArray
-APyFloatArray::squeeze(std::optional<std::variant<nb::int_, nb::tuple>> axis) const
-{
-    std::vector<std::size_t> shape = _shape;
-    std::set<int> axis_set;
-    if (axis.has_value()) {
-        auto ax = axis.value();
-        nb::tuple axis_tuple;
-
-        if (std::holds_alternative<nb::tuple>(ax)) {
-            axis_tuple = std::get<nb::tuple>(ax);
-        } else if (std::holds_alternative<nb::int_>(ax)) {
-            int axis = int(std::get<nb::int_>(ax));
-            axis_tuple = nb::make_tuple(axis);
-        }
-        for (auto ptr = axis_tuple.begin(); ptr != axis_tuple.end(); ptr++) {
-            int axis_n = int(nanobind::cast<nb::int_>(*ptr));
-            if (axis_n >= int(shape.size())) {
-                throw nb::index_error(
-                    "Specified axis with larger than number of dimensions in the "
-                    "APyFloatArray"
-                );
-            }
-            axis_set.insert(axis_n);
-        }
-        // precidate to squeeze without specified axes
-        int cnt = 0;
-        auto predicate = [&](std::size_t dim) {
-            if (axis_set.find(cnt) != axis_set.end() && dim != 1) {
-                throw nb::value_error(
-                    "Cannot select an axis to squeeze that has size other than one"
-                );
-            }
-            return axis_set.find(cnt++) != axis_set.end() && dim == 1;
-        };
-        shape.erase(std::remove_if(shape.begin(), shape.end(), predicate), shape.end());
-    } else {
-        // precidate to squeeze without specified axes
-        auto predicate = [](std::size_t dim) { return dim == 1; };
-        shape.erase(std::remove_if(shape.begin(), shape.end(), predicate), shape.end());
-    }
-
-    if (shape.size() == 0) {
-        shape = { 1 };
-    }
-
-    // create the resulting Array and copy the data
-    APyFloatArray result(shape, exp_bits, man_bits, bias);
-    std::copy_n(_data.begin(), _data.size(), result._data.begin());
-    return result;
-}
-
 //! Perform a linear convolution with `other` using `mode`
 APyFloatArray
 APyFloatArray::convolve(const APyFloatArray& other, const std::string& mode) const
@@ -1790,18 +1734,6 @@ std::string APyFloatArray::repr() const
     return ss.str();
 }
 
-APyFloatArray APyFloatArray::flatten() const
-{
-    // Reuse the reshape function to flatten the array
-    return this->reshape(nb::make_tuple(-1));
-}
-
-APyFloatArray APyFloatArray::ravel() const
-{
-    // currently same as flatten
-    return this->flatten();
-}
-
 std::variant<
     nb::list,
     nb::ndarray<nb::numpy, uint64_t>,
@@ -2116,55 +2048,6 @@ APyFloatArray APyFloatArray::swapaxes(nb::int_ axis1, nb::int_ axis2) const
     return result;
 }
 
-APyFloatArray APyFloatArray::transpose(std::optional<nb::tuple> axes) const
-{
-    switch (ndim()) {
-    case 0:
-    case 1:
-        // Behave like `NumPy`, simply returns `*this` if single-dimensional
-        return *this;
-    case 2: {
-
-        // Optimized code for dim == 2
-        // Resulting array with shape dimensions
-        std::vector<size_t> new_shape = { _shape[1], _shape[0] };
-        APyFloatArray result(new_shape, exp_bits, man_bits, bias);
-
-        // Transpose the data
-        for (std::size_t y = 0; y < _shape[0]; ++y) {
-            for (std::size_t x = 0; x < _shape[1]; ++x) {
-                result._data[x * _shape[0] + y] = _data[y * _shape[1] + x];
-            }
-        }
-        return result;
-    }
-
-    default: {
-        std::vector<size_t> new_axis(ndim());
-
-        if (axes.has_value()) {
-            std::variant<nb::tuple, nb::int_> axis_variant = axes.value();
-            new_axis = get_normalized_axes(axis_variant, ndim());
-        } else {
-            // reverse the order of axes by default
-            std::iota(new_axis.begin(), new_axis.end(), 0);
-            std::reverse(new_axis.begin(), new_axis.end());
-        }
-
-        std::vector<size_t> new_shape(ndim());
-        for (size_t i = 0; i < ndim(); ++i) {
-            new_shape[i] = _shape[new_axis[i]];
-        }
-
-        APyFloatArray result(new_shape, exp_bits, man_bits, bias);
-        transpose_axes_and_copy_data(
-            _data.begin(), result._data.begin(), _shape, new_axis
-        );
-        return result;
-    }
-    }
-}
-
 APyFloatArray APyFloatArray::cast(
     std::optional<int> new_exp_bits,
     std::optional<int> new_man_bits,
diff --git a/src/apyfloatarray.h b/src/apyfloatarray.h
index e6701c6c3..755e88c3a 100644
--- a/src/apyfloatarray.h
+++ b/src/apyfloatarray.h
@@ -179,16 +179,6 @@ class APyFloatArray : public APyArray<APyFloatData, APyFloatArray> {
      * *                          Public member functions                           * *
      * ****************************************************************************** */
 
-    //! Reshape flatten into 1d shape.
-    APyFloatArray flatten() const;
-
-    //! Same as flatten as for now.
-    APyFloatArray ravel() const;
-
-    //! Return a copy where all axes of size 1 is removed.
-    APyFloatArray
-    squeeze(std::optional<std::variant<nb::int_, nb::tuple>> axis = std::nullopt) const;
-
     //! Perform a linear convolution with `other` using `mode`
     APyFloatArray convolve(const APyFloatArray& other, const std::string& mode) const;
 
@@ -239,10 +229,6 @@ class APyFloatArray : public APyArray<APyFloatData, APyFloatArray> {
     //! Swaps the positions of two axes in the array.
     APyFloatArray swapaxes(nb::int_ axis1, nb::int_ axis2) const;
 
-    //! Transposition function. For a 1-D array, Return an exact copy of `*this`. For
-    //! a 2-D array, Return the matrix transposition of `*this`.
-    APyFloatArray transpose(std::optional<nb::tuple> axes = std::nullopt) const;
-
     //! Return a copy of the tensor with the elements resized.
     APyFloatArray cast(
         std::optional<int> exp_bits,
diff --git a/src/array_utils.h b/src/array_utils.h
index 918a8f97e..5dcf35de8 100644
--- a/src/array_utils.h
+++ b/src/array_utils.h
@@ -6,12 +6,14 @@
 #include "apytypes_util.h"
 #include "broadcast.h"
 #include "nanobind/ndarray.h"
-#include <cstddef>
+
 #include <nanobind/nanobind.h>
-#include <unordered_set>
-#include <variant>
+#include <nanobind/stl/variant.h>
 namespace nb = nanobind;
 
+#include <cstddef>
+#include <unordered_set>
+#include <variant>
 #include <vector>
 
 /**


Test patch:
diff --git a/lib/test/apyfixedarray/test_methods.py b/lib/test/apyfixedarray/test_methods.py
index 51c71bc23..64100b205 100644
--- a/lib/test/apyfixedarray/test_methods.py
+++ b/lib/test/apyfixedarray/test_methods.py
@@ -42,15 +42,14 @@ def test_squeeze():
     assert z.is_identical(i)
     m = APyFixedArray([[1], [2], [3]], bits=2, int_bits=2)
     with pytest.raises(
-        ValueError,
-        match="Cannot select an axis to squeeze that has size other than one",
+        ValueError, match=r"APyFixedArray\.squeeze: cannot squeeze non-one dimensions"
     ):
         _ = m.squeeze(axis=0)
     m1 = m.squeeze(axis=1)
     assert m1.is_identical(APyFixedArray([1, 2, 3], bits=2, int_bits=2))
     with pytest.raises(
         IndexError,
-        match="Specified axis with larger than number of dimensions in the APyFixedArray",
+        match=r"APyFixedArray\.squeeze: specified axis larger than number of array dim",
     ):
         _ = m.squeeze(axis=2)
 
@@ -61,13 +60,12 @@ def test_squeeze():
     assert p.is_identical(APyFixedArray([[[[[2]]]]], bits=2, int_bits=2))
     q = APyFixedArray([[[1]], [[2]], [[3]], [[4]]], bits=2, int_bits=2)
     with pytest.raises(
-        ValueError,
-        match="Cannot select an axis to squeeze that has size other than one",
+        ValueError, match=r"APyFixedArray\.squeeze: cannot squeeze non-one dimensions"
     ):
         _ = q.squeeze((0, 1, 2))
     with pytest.raises(
         IndexError,
-        match="Specified axis with larger than number of dimensions in the APyFixedArray",
+        match=r"APyFixedArray.squeeze: specified axis larger than number of array dim",
     ):
         _ = m.squeeze((1, 4))