Repository: apytypes/apytypes (language: python)

Issue:
`__setitem__`
It will be useful if one can set parts of an array to either a scalar or an array (of the same type/word length).


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7780a04b9..631d8a12b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added optional parameter for `bias` in array creation functions that were
   missing it.
 - Added `APyFloatArray.from_bits`.
+- Added `APyFixedArray.__setitem__` and `APyFloatArray.__setitem__`.
 
 ### Fixed
 
diff --git a/lib/apytypes/_apytypes.pyi b/lib/apytypes/_apytypes.pyi
index 64d877dcf..2f24889c7 100644
--- a/lib/apytypes/_apytypes.pyi
+++ b/lib/apytypes/_apytypes.pyi
@@ -1486,8 +1486,13 @@ class APyFixedArray:
     def __rshift__(self, shift_amnt: int) -> APyFixedArray: ...
     def __abs__(self) -> APyFixedArray: ...
     def __getitem__(
-        self, idx: int | slice | types.EllipsisType | tuple
+        self, key: int | slice | types.EllipsisType | tuple
     ) -> APyFixedArray | APyFixed: ...
+    def __setitem__(
+        self,
+        key: int | slice | types.EllipsisType | tuple,
+        val: APyFixedArray | APyFixed,
+    ) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> APyFixedArrayIterator: ...
     def __array__(self) -> Annotated[ArrayLike, dict(dtype="float64")]: ...
@@ -3267,8 +3272,13 @@ class APyFloatArray:
         """
 
     def __getitem__(
-        self, idx: int | slice | types.EllipsisType | tuple
+        self, key: int | slice | types.EllipsisType | tuple
     ) -> APyFloatArray | APyFloat: ...
+    def __setitem__(
+        self,
+        key: int | slice | types.EllipsisType | tuple,
+        val: APyFloatArray | APyFloat,
+    ) -> None: ...
     def __iter__(self) -> APyFloatArrayIterator: ...
     def __array__(self) -> Annotated[ArrayLike, dict(dtype="float64")]: ...
     def cast(
diff --git a/src/apyarray.h b/src/apyarray.h
index 4d63e0c63..c3b50c1fc 100644
--- a/src/apyarray.h
+++ b/src/apyarray.h
@@ -7,19 +7,23 @@
 #define _APYARRAY_H
 
 #include "apybuffer.h"
+#include "apytypes_fwd.h"
 #include "apytypes_util.h"
 #include "array_utils.h"
 #include "broadcast.h"
+#include "fmt/format.h"
 
 // Python object access through Nanobind
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/variant.h> // std::variant (with nanobind support)
 namespace nb = nanobind;
 
-#include <cassert>  // assert
-#include <iterator> // std::begin
-#include <set>      // std::set
-#include <variant>  // std::variant
+#include <cassert>     // assert
+#include <exception>   // std::exception
+#include <iterator>    // std::begin
+#include <set>         // std::set
+#include <string_view> // std::string_view
+#include <variant>     // std::variant
 
 template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T> {
 
@@ -43,13 +47,16 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
     //! Convert a possibly negative integer index `idx` to a positive one if used to
     //! index dimension `dim` of `*this`. Throws `std::out_of_range` if the index could
     //! not be converted or is out of range.
-    std::size_t adjust_integer_index(std::ptrdiff_t idx, std::size_t dim) const
+    std::size_t adjust_integer_index(
+        std::ptrdiff_t idx, std::size_t dim, std::string_view func_name
+    ) const
     {
         // Bounds checking
         if (idx >= std::ptrdiff_t(_shape[dim]) || idx < -std::ptrdiff_t(_shape[dim])) {
             std::string msg = fmt::format(
-                "{}.__getitem__: index {} is out of bounds for axis {} with size {}",
+                "{}.{}: index {} is out of bounds for axis {} with size {}",
                 ARRAY_TYPE::ARRAY_NAME,
+                func_name,
                 idx,
                 dim,
                 _shape[dim]
@@ -60,28 +67,27 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
         return idx < 0 ? idx + _shape[dim] : idx;
     }
 
-    //! Compute the resulting shape when slicing `*this` with non-empty C++ `tuple`.
-    //! Calling this function requires that `tuple.size() <= _shape.size()`
+    //! Compute the resulting shape when slicing `*this` with non-empty C++ tuple `key`.
+    //! Calling this function requires that `key.size() <= _shape.size()`
     std::vector<std::size_t>
-    compute_slice_shape(const std::vector<std::variant<nb::int_, nb::slice>>& tuple
-    ) const
+    compute_slice_shape(const std::vector<std::variant<nb::int_, nb::slice>>& key) const
     {
-        assert(tuple.size() <= _shape.size());
+        assert(key.size() <= _shape.size());
         std::vector<std::size_t> shape;
 
-        // Elements in the tuple
-        for (std::size_t i = 0; i < tuple.size(); i++) {
-            if (std::holds_alternative<nb::int_>(tuple[i])) {
+        // Elements in the key tuple
+        for (std::size_t i = 0; i < key.size(); i++) {
+            if (std::holds_alternative<nb::int_>(key[i])) {
                 continue;
-            } else { /* std::hold_alternative<nb::slice>(tuple[i]) */
-                auto&& slice = std::get<nb::slice>(tuple[i]);
+            } else { /* std::hold_alternative<nb::slice>(key[i]) */
+                auto&& slice = std::get<nb::slice>(key[i]);
                 auto [_, __, ___, len] = slice.compute(_shape[i]);
                 shape.push_back(len);
             }
         }
 
-        // Any dimension in `*this` not referred to in `tuple` simply "come along"
-        for (std::size_t i = tuple.size(); i < _shape.size(); i++) {
+        // Any dimension in `*this` not referred to in `key` simply "come along"
+        for (std::size_t i = key.size(); i < _shape.size(); i++) {
             shape.push_back(_shape[i]);
         }
 
@@ -91,14 +97,16 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
     //! Convert a `nb::tuple` of `nb::int_`, `nb::slice`, and `nb::ellipsis` to
     //! a `std::vector<std::variant<nb::int_, nb::slice>>`, resolving any ellipsis
     //! slicing.
-    auto resolve_python_tuple_slice(const nb::tuple& key) const
+    std::vector<std::variant<nb::int_, nb::slice>>
+    resolve_python_tuple_slice(const nb::tuple& key, std::string_view func_name) const
     {
         if (key.size() > _ndim) {
             // The key tuple must have fewer elements than this array has number of
             // dimensions
             std::string msg = fmt::format(
-                "{}.__getitem__: key tuple size (={}) > ndim (={})",
+                "{}.{}: key tuple size (={}) > ndim (={})",
                 ARRAY_TYPE::ARRAY_NAME,
+                func_name,
                 key.size(),
                 _ndim
             );
@@ -122,8 +130,9 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
                     // An ellipsis object has already been procecced. Only a single
                     // ellipsis per tuple key is allowed
                     std::string msg = fmt::format(
-                        "{}.__getitem__: only one single ellipsis object allowed",
-                        ARRAY_TYPE::ARRAY_NAME
+                        "{}.{}: only one single ellipsis object allowed",
+                        ARRAY_TYPE::ARRAY_NAME,
+                        func_name
                     );
                     throw nb::value_error(msg.c_str());
                 } else {
@@ -136,9 +145,10 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
                 }
             } else {
                 std::string msg = fmt::format(
-                    "{}.__getitem__: supported keys are `int`, `slice`, `ellipsis`, or "
-                    "a single tuple thereof",
-                    ARRAY_TYPE::ARRAY_NAME
+                    "{}.{}: supported keys are `int`, `slice`, `ellipsis`, or a single "
+                    "tuple thereof",
+                    ARRAY_TYPE::ARRAY_NAME,
+                    func_name
                 );
                 throw nb::value_error(msg.c_str());
             }
@@ -147,13 +157,14 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
     }
 
     //! Retrieve item from integer index
-    auto get_item_integer(std::ptrdiff_t idx) const
+    std::variant<ARRAY_TYPE, scalar_variant_t<ARRAY_TYPE>>
+    get_item_integer(std::ptrdiff_t idx) const
     {
-        using SCALAR_TYPE = typename ARRAY_TYPE::SCALAR_VARIANT;
+        using SCALAR_TYPE = scalar_variant_t<ARRAY_TYPE>;
         using RESULT_TYPE = std::variant<ARRAY_TYPE, SCALAR_TYPE>;
 
         // Adjust a possibly negative index
-        idx = adjust_integer_index(idx, 0);
+        idx = adjust_integer_index(idx, 0, "__getitem__");
 
         if (_ndim == 1) {
             SCALAR_TYPE result = static_cast<const ARRAY_TYPE*>(this)->create_scalar();
@@ -197,7 +208,7 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
              */
             if (std::holds_alternative<nb::int_>(element)) {
                 auto idx = static_cast<std::ptrdiff_t>(std::get<nb::int_>(element));
-                idx = adjust_integer_index(idx, dim);
+                idx = adjust_integer_index(idx, dim, "__getitem__");
                 std::copy_n(
                     input_it + _itemsize * idx * strides[dim],
                     _itemsize * strides[dim],
@@ -236,7 +247,7 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
              */
             if (std::holds_alternative<nb::int_>(element)) {
                 auto idx = static_cast<std::ptrdiff_t>(std::get<nb::int_>(element));
-                idx = adjust_integer_index(idx, dim);
+                idx = adjust_integer_index(idx, dim, "__getitem__");
                 return get_item_tuple_recursive_descent(
                     tuple,
                     strides,
@@ -277,11 +288,11 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
 
     //! Retrieve item(s) from a `std::vector` of `nb::int_` and `nb::slice`. Assumes
     //! that `tuple.size() <= _shape.size()`.
-    auto get_item_tuple(const std::vector<std::variant<nb::int_, nb::slice>>& tuple
-    ) const
+    std::variant<ARRAY_TYPE, scalar_variant_t<ARRAY_TYPE>>
+    get_item_tuple(const std::vector<std::variant<nb::int_, nb::slice>>& tuple) const
     {
         assert(tuple.size() <= _shape.size());
-        using SCALAR_TYPE = typename ARRAY_TYPE::SCALAR_VARIANT;
+        using SCALAR_TYPE = scalar_variant_t<ARRAY_TYPE>;
         using RESULT_TYPE = std::variant<ARRAY_TYPE, SCALAR_TYPE>;
 
         // Return everything on an empty tuple
@@ -305,7 +316,7 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
             for (std::size_t i = 0; i < tuple.size(); i++) {
                 assert(std::holds_alternative<nb::int_>(tuple[i]));
                 auto axis = static_cast<std::ptrdiff_t>(std::get<nb::int_>(tuple[i]));
-                axis = adjust_integer_index(axis, i);
+                axis = adjust_integer_index(axis, i, "__getitem__");
                 item_idx += strides[i] * axis;
             }
             result.copy_n_from(std::begin(_data) + item_idx * _itemsize, _itemsize);
@@ -327,9 +338,11 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
     }
 
     //! Top-level Python exported `__getitem__` function
-    auto get_item(std::variant<nb::int_, nb::slice, nb::ellipsis, nb::tuple> key) const
+    std::variant<ARRAY_TYPE, scalar_variant_t<ARRAY_TYPE>>
+    get_item(const std::variant<nb::int_, nb::slice, nb::ellipsis, nb::tuple>& key
+    ) const
     {
-        using SCALAR_TYPE = typename ARRAY_TYPE::SCALAR_VARIANT;
+        using SCALAR_TYPE = scalar_variant_t<ARRAY_TYPE>;
         using RESULT_TYPE = std::variant<ARRAY_TYPE, SCALAR_TYPE>;
 
         if (std::holds_alternative<nb::int_>(key)) {
@@ -343,7 +356,9 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
 
             // Key is of slice type
             nb::tuple tuple = nb::make_tuple(std::get<nb::slice>(key));
-            return RESULT_TYPE(get_item_tuple(resolve_python_tuple_slice(tuple)));
+            return RESULT_TYPE(
+                get_item_tuple(resolve_python_tuple_slice(tuple, "__getitem__"))
+            );
 
         } else if (std::holds_alternative<nb::ellipsis>(key)) {
 
@@ -353,9 +368,9 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
         } else { /* std::holds_alternative<nb::tuple>(key) */
 
             // Key is a tuple of slicing instructions
-            return RESULT_TYPE(
-                get_item_tuple(resolve_python_tuple_slice(std::get<nb::tuple>(key)))
-            );
+            return RESULT_TYPE(get_item_tuple(
+                resolve_python_tuple_slice(std::get<nb::tuple>(key), "__getitem__")
+            ));
         }
     }
 
@@ -363,6 +378,219 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
      * *                    `__setitem__` family of methods                         * *
      * ****************************************************************************** */
 
+    //! Working part of `set_item`, using recursive descent to copy items from
+    //! `input_it` to `output_it`. Returns the number of elements copied. Assumes that
+    //! `dim < key.size()`.
+    std::size_t set_item_recursive_descent(
+        const std::vector<std::variant<nb::int_, nb::slice>>& key,
+        const std::vector<std::size_t>& strides,
+        typename APyBuffer<T>::vector_type::const_iterator input_it,
+        typename APyBuffer<T>::vector_type::iterator output_it,
+        std::size_t dim = 0
+    )
+    {
+        auto&& element = key[dim];
+        if (dim == key.size() - 1) {
+            /*
+             * Final dimension in the tuple key of the recursive descent. Do the
+             * actual copying of data.
+             */
+            if (std::holds_alternative<nb::int_>(element)) {
+                auto idx = static_cast<std::ptrdiff_t>(std::get<nb::int_>(element));
+                idx = adjust_integer_index(idx, dim, "__setitem__");
+                std::copy_n(
+                    input_it,
+                    _itemsize * strides[dim],
+                    output_it + _itemsize * strides[dim] * idx
+                );
+                return strides[dim];
+            } else { /* std::holds_alternative<nb::slice>(element) */
+                auto&& slice = std::get<nb::slice>(element);
+                auto [start, stop, step, len] = slice.compute(_shape[dim]);
+                std::size_t elements_copied = 0;
+                if (step < 0) {
+                    for (std::ptrdiff_t src_i = start; src_i > stop; src_i += step) {
+                        std::copy_n(
+                            input_it + _itemsize * elements_copied,
+                            _itemsize * strides[dim],
+                            output_it + _itemsize * src_i * strides[dim]
+                        );
+                        elements_copied += strides[dim];
+                    }
+                } else { /* step >= 0 */
+                    for (std::ptrdiff_t src_i = start; src_i < stop; src_i += step) {
+                        std::copy_n(
+                            input_it + _itemsize * elements_copied,
+                            _itemsize * strides[dim],
+                            output_it + _itemsize * src_i * strides[dim]
+                        );
+                        elements_copied += strides[dim];
+                    }
+                }
+                return elements_copied;
+            }
+        } else {
+            /*
+             * This is not yet the final dimesnsion of the tuple slice in this recursive
+             * descent. We need to go deeper...
+             */
+            if (std::holds_alternative<nb::int_>(element)) {
+                auto idx = static_cast<std::ptrdiff_t>(std::get<nb::int_>(element));
+                idx = adjust_integer_index(idx, dim, "__setitem__");
+                return set_item_recursive_descent(
+                    key,
+                    strides,
+                    input_it,
+                    output_it + _itemsize * idx * strides[dim],
+                    dim + 1
+                );
+            } else { /* std::holds_alternative<nb::slice>(tuple_element) */
+                auto&& slice = std::get<nb::slice>(element);
+                auto [start, stop, step, len] = slice.compute(_shape[dim]);
+                std::size_t elements_copied = 0;
+                if (step < 0) {
+                    for (std::ptrdiff_t src_i = start; src_i > stop; src_i += step) {
+                        elements_copied += set_item_recursive_descent(
+                            key,
+                            strides,
+                            input_it + _itemsize * elements_copied,
+                            output_it + _itemsize * src_i * strides[dim],
+                            dim + 1
+                        );
+                    }
+                } else { /* step >= 0 */
+                    for (std::ptrdiff_t src_i = start; src_i < stop; src_i += step) {
+                        elements_copied += set_item_recursive_descent(
+                            key,
+                            strides,
+                            input_it + _itemsize * elements_copied,
+                            output_it + _itemsize * src_i * strides[dim],
+                            dim + 1
+                        );
+                    }
+                }
+
+                return elements_copied;
+            }
+        }
+    }
+
+    void set_item_from_array(
+        const std::vector<std::variant<nb::int_, nb::slice>>& key, const ARRAY_TYPE& val
+    )
+    {
+        assert(key.size() <= _shape.size());
+
+        // Make sure that all bit specifiers in `*this` and `val` are equal.
+        if (!static_cast<const ARRAY_TYPE*>(this)->same_type_as(val)) {
+            std::string error_msg = fmt::format(
+                "{}.__setitem__: `val` has different bit specifiers than `self`",
+                ARRAY_TYPE::ARRAY_NAME
+            );
+            throw nb::value_error(error_msg.c_str());
+        }
+
+        // Compute the slice shape
+        std::vector<std::size_t> slice_shape = compute_slice_shape(key);
+        if (!slice_shape.size()) {
+            // `slice_shape` is empty *iff* integers are used to slice all dimensions
+            slice_shape.push_back(1);
+        }
+
+        // Compute the stride of `*this`
+        std::vector<std::size_t> strides = strides_from_shape(_shape);
+
+        if (val._shape == slice_shape) {
+            /*
+             * Shape of value array and `*this` slice matches up exactly, simply copy
+             * the data using recursive descent.
+             */
+            set_item_recursive_descent(
+                key, strides, std::begin(val._data), std::begin(_data)
+            );
+        } else if (is_broadcastable(val._shape, slice_shape)) {
+            /*
+             * The value shape can be broadcast to the slice shape. Broadcast first, and
+             * then copy the data using recursive descent.
+             */
+            ARRAY_TYPE val_broadcasted
+                = static_cast<const ARRAY_TYPE*>(this)->create_array(slice_shape);
+            broadcast_data_copy(
+                std::begin(val._data),             // src
+                std::begin(val_broadcasted._data), // dst
+                val._shape,                        // src_shape
+                slice_shape,                       // dst_shape
+                _itemsize                          // itemsize
+            );
+            set_item_recursive_descent(
+                key, strides, std::begin(val_broadcasted._data), std::begin(_data)
+            );
+        } else {
+            std::string error_msg = fmt::format(
+                "{}.__setitem__: `val` shape not broadcastable to `key` slice",
+                ARRAY_TYPE::ARRAY_NAME
+            );
+            throw nb::value_error(error_msg.c_str());
+        }
+    }
+
+    void set_item_from_scalar(
+        const std::vector<std::variant<nb::int_, nb::slice>>& key,
+        const scalar_variant_t<ARRAY_TYPE>& val
+    )
+    {
+        // Make sure that all bit specifiers in `*this` and `val` are equal.
+        if (!static_cast<const ARRAY_TYPE*>(this)->same_type_as(val)) {
+            std::string error_msg = fmt::format(
+                "{}.__setitem__: `val` has different bit specifiers than `self`",
+                ARRAY_TYPE::ARRAY_NAME
+            );
+            throw nb::value_error(error_msg.c_str());
+        }
+
+        // Make an array out of the scalar and set item from there
+        ARRAY_TYPE array_val
+            = static_cast<const ARRAY_TYPE*>(this)->create_array({ 1 });
+        val.copy_n_to(std::begin(array_val._data), _itemsize);
+        set_item_from_array(key, array_val);
+    }
+
+    //! Python exported `__setitem__` method for APyArrays
+    void set_item(
+        const std::variant<nb::int_, nb::slice, nb::ellipsis, nb::tuple>& key,
+        const std::variant<std::monostate, ARRAY_TYPE, scalar_variant_t<ARRAY_TYPE>>&
+            val
+    )
+    {
+        // Convert the variant of Python objects into a tuple of keys
+        nb::tuple python_tuple_key;
+        if (std::holds_alternative<nb::int_>(key)) {
+            python_tuple_key = nb::make_tuple(std::get<nb::int_>(key));
+        } else if (std::holds_alternative<nb::slice>(key)) {
+            python_tuple_key = nb::make_tuple(std::get<nb::slice>(key));
+        } else if (std::holds_alternative<nb::ellipsis>(key)) {
+            python_tuple_key = nb::make_tuple(std::get<nb::ellipsis>(key));
+        } else { /* std::holds_alternative<nb::ellipsis>(key) */
+            python_tuple_key = std::get<nb::tuple>(key);
+        }
+
+        // Resolve the tuple of ellipsis objects and flatten to `std::vector` key
+        std::vector<std::variant<nb::int_, nb::slice>> cpp_key
+            = resolve_python_tuple_slice(python_tuple_key, "__setitem__");
+
+        if (std::holds_alternative<scalar_variant_t<ARRAY_TYPE>>(val)) {
+            set_item_from_scalar(cpp_key, std::get<scalar_variant_t<ARRAY_TYPE>>(val));
+        } else if (std::holds_alternative<ARRAY_TYPE>(val)) {
+            set_item_from_array(cpp_key, std::get<ARRAY_TYPE>(val));
+        } else {
+            std::string msg = fmt::format(
+                "{}.__setitem__: unsupported value type (internal error: `monostate`?)",
+                ARRAY_TYPE::ARRAY_NAME
+            );
+            throw nb::value_error(msg.c_str());
+        }
+    }
+
     /* ****************************************************************************** *
      * *                     `broadcast` family of methods                          * *
      * ****************************************************************************** */
@@ -465,7 +693,7 @@ template <typename T, typename ARRAY_TYPE> class APyArray : public APyBuffer<T>
         std::vector<std::size_t> cpp_shape = try_reshape(shape);
         ARRAY_TYPE result
             = static_cast<const ARRAY_TYPE*>(this)->create_array(cpp_shape);
-        std::copy_n(this->_data.begin(), _data.size(), result._data.begin());
+        std::copy_n(_data.begin(), _data.size(), result._data.begin());
 
         return result;
     }
diff --git a/src/apyfixed.h b/src/apyfixed.h
index 86045786c..831de4c33 100644
--- a/src/apyfixed.h
+++ b/src/apyfixed.h
@@ -66,18 +66,26 @@ class APyFixed {
      * ****************************************************************************** */
 
 public:
+    //! Copy `n` items from `it` into `*this`
     template <typename RANDOM_ACCESS_ITERATOR>
     void copy_n_from(RANDOM_ACCESS_ITERATOR it, std::size_t n) noexcept
     {
         std::copy_n(it, n, std::begin(_data));
     }
 
+    //! Copy `n` items from `*this` into `it`
     template <typename RANDOM_ACCESS_ITERATOR>
     void copy_n_to(RANDOM_ACCESS_ITERATOR it, std::size_t n) const noexcept
     {
         std::copy_n(std::begin(_data), n, it);
     }
 
+    //! Test if two fixed-point numbers have the same bit specifiers
+    APY_INLINE bool same_type_as(const APyFixed& other) const
+    {
+        return _bits == other._bits && _int_bits == other._int_bits;
+    }
+
     /* ****************************************************************************** *
      *                            Python constructors                                 *
      * ****************************************************************************** */
diff --git a/src/apyfixedarray.h b/src/apyfixedarray.h
index b315f1110..1c00dc88f 100644
--- a/src/apyfixedarray.h
+++ b/src/apyfixedarray.h
@@ -65,9 +65,6 @@ class APyFixedArray : public APyArray<mp_limb_t, APyFixedArray> {
      * ****************************************************************************** */
 
 public:
-    //! The scalar variant of `APyFixedArray` is `APyFixed`
-    using SCALAR_VARIANT = APyFixed;
-
     //! Name of this array type (used when throwing errors)
     static constexpr auto ARRAY_NAME = std::string_view("APyFixedArray");
 
@@ -78,6 +75,18 @@ class APyFixedArray : public APyArray<mp_limb_t, APyFixedArray> {
         return APyFixedArray(shape, _bits, _int_bits);
     }
 
+    //! Test if two fixed-point vectors have the same bit specifiers
+    APY_INLINE bool same_type_as(const APyFixedArray& other) const
+    {
+        return _bits == other._bits && _int_bits == other._int_bits;
+    }
+
+    //! Test if `*this` has the same bit specifiers as another `APyFixed`
+    APY_INLINE bool same_type_as(const APyFixed& other) const
+    {
+        return _bits == other._bits && _int_bits == other._int_bits;
+    }
+
     /* ****************************************************************************** *
      * *                          Python constructors                               * *
      * ****************************************************************************** */
diff --git a/src/apyfixedarray_wrapper.cc b/src/apyfixedarray_wrapper.cc
index a9daebcd6..de345f837 100644
--- a/src/apyfixedarray_wrapper.cc
+++ b/src/apyfixedarray_wrapper.cc
@@ -1279,7 +1279,8 @@ void bind_fixed_array(nb::module_& m)
         .def("__abs__", &APyFixedArray::abs)
 
         // Iteration and friends
-        .def("__getitem__", &APyFixedArray::get_item, nb::arg("idx"))
+        .def("__getitem__", &APyFixedArray::get_item, nb::arg("key"))
+        .def("__setitem__", &APyFixedArray::set_item, nb::arg("key"), nb::arg("val"))
         .def("__len__", &APyFixedArray::size)
         .def(
             "__iter__",
diff --git a/src/apyfloat.h b/src/apyfloat.h
index c7b6ec4af..fd6ad6922 100644
--- a/src/apyfloat.h
+++ b/src/apyfloat.h
@@ -420,6 +420,7 @@ class APyFloat {
      * ****************************************************************************** */
 
 public:
+    //! Copy `n` items from `it` into `*this`
     template <typename RANDOM_ACCESS_ITERATOR>
     void copy_n_from(RANDOM_ACCESS_ITERATOR it, std::size_t n) noexcept
     {
@@ -429,6 +430,7 @@ class APyFloat {
         exp = it->exp;
     }
 
+    //! Copy `n` items from `*this` into `it`
     template <typename RANDOM_ACCESS_ITERATOR>
     void copy_n_to(RANDOM_ACCESS_ITERATOR it, std::size_t n) const noexcept
     {
@@ -438,6 +440,13 @@ class APyFloat {
         it->exp = exp;
     }
 
+    //! Test if two floating-point numbers have the same bit specifiers
+    APY_INLINE bool same_type_as(const APyFloat& other) const
+    {
+        return man_bits == other.man_bits && exp_bits == other.exp_bits
+            && bias == other.bias;
+    }
+
     /* ******************************************************************************
      * * Non-Python accessible constructors                                         *
      * ******************************************************************************
@@ -468,14 +477,6 @@ class APyFloat {
     //! Create a floating-point object that is normalized
     APyFloat normalized() const;
 
-    //! Test if two floating-point numbers are identical, i.e., has the same value, and
-    //! the same format
-    APY_INLINE bool same_type_as(const APyFloat& other) const
-    {
-        return man_bits == other.man_bits && exp_bits == other.exp_bits
-            && bias == other.bias;
-    }
-
     template <typename T, typename ARRAY_TYPE> friend class APyArray;
 };
 
diff --git a/src/apyfloatarray.h b/src/apyfloatarray.h
index 49f4f0da2..d9273d271 100644
--- a/src/apyfloatarray.h
+++ b/src/apyfloatarray.h
@@ -39,9 +39,6 @@ class APyFloatArray : public APyArray<APyFloatData, APyFloatArray> {
      * ****************************************************************************** */
 
 public:
-    //! The scalar variant of `APyFloatArray` is `APyFloat`
-    using SCALAR_VARIANT = APyFloat;
-
     //! Name of this array type (used when throwing errors)
     static constexpr auto ARRAY_NAME = std::string_view("APyFloatArray");
 
@@ -52,6 +49,20 @@ class APyFloatArray : public APyArray<APyFloatData, APyFloatArray> {
         return APyFloatArray(shape, exp_bits, man_bits, bias);
     }
 
+    //! Test if two floating-point vectors have the same bit specifiers
+    APY_INLINE bool same_type_as(const APyFloatArray& other) const
+    {
+        return man_bits == other.man_bits && exp_bits == other.exp_bits
+            && bias == other.bias;
+    }
+
+    //! Test if `*this` has the same bit specifiers as another `APyFloat`
+    APY_INLINE bool same_type_as(const APyFloat& other) const
+    {
+        return man_bits == other.get_man_bits() && exp_bits == other.get_exp_bits()
+            && bias == other.get_bias();
+    }
+
     /* ****************************************************************************** *
      * *                       Binary arithmetic operators                          * *
      * ****************************************************************************** */
@@ -275,20 +286,6 @@ class APyFloatArray : public APyArray<APyFloatData, APyFloatArray> {
      */
     bool is_identical(const APyFloatArray& other) const;
 
-    //! Test if two `APyFloatArray` objects have the same format
-    APY_INLINE bool same_type_as(const APyFloatArray& other) const
-    {
-        return man_bits == other.man_bits && exp_bits == other.exp_bits
-            && bias == other.bias;
-    }
-
-    //! Test if the `APyFloatArray` object has the same format as a `APyFloat` object
-    APY_INLINE bool same_type_as(const APyFloat& other) const
-    {
-        return man_bits == other.get_man_bits() && exp_bits == other.get_exp_bits()
-            && bias == other.get_bias();
-    }
-
     /* ****************************************************************************** *
      * *            `__getitem__` and `__setitem__` family of methods               * *
      * ****************************************************************************** */
diff --git a/src/apyfloatarray_wrapper.cc b/src/apyfloatarray_wrapper.cc
index d80506dbf..eb3d82385 100644
--- a/src/apyfloatarray_wrapper.cc
+++ b/src/apyfloatarray_wrapper.cc
@@ -1380,7 +1380,8 @@ void bind_float_array(nb::module_& m)
         )
 
         // Iteration and friends
-        .def("__getitem__", &APyFloatArray::get_item, nb::arg("idx"))
+        .def("__getitem__", &APyFloatArray::get_item, nb::arg("key"))
+        .def("__setitem__", &APyFloatArray::set_item, nb::arg("key"), nb::arg("val"))
         .def(
             "__iter__",
             [](nb::iterable array) {
diff --git a/src/apytypes_fwd.h b/src/apytypes_fwd.h
new file mode 100644
index 000000000..917be1377
--- /dev/null
+++ b/src/apytypes_fwd.h
@@ -0,0 +1,35 @@
+/*
+ * Forward declarations of the important APyTypes types, and their relationship
+ */
+
+// Scalar types
+class APyFixed;
+class APyFloat;
+
+// Array types
+class APyFixedArray;
+class APyFloatArray;
+
+/*
+ * APyTypes relationships
+ */
+template <typename T> struct scalar_variant;
+
+template <typename T> struct array_variant;
+
+template <> struct scalar_variant<APyFixedArray> {
+    using type = APyFixed;
+};
+template <> struct scalar_variant<APyFloatArray> {
+    using type = APyFloat;
+};
+template <> struct array_variant<APyFixed> {
+    using type = APyFixedArray;
+};
+template <> struct array_variant<APyFloat> {
+    using type = APyFloatArray;
+};
+
+template <typename T> using scalar_variant_t = typename scalar_variant<T>::type;
+
+template <typename T> using array_variant_t = typename array_variant<T>::type;
diff --git a/src/broadcast.h b/src/broadcast.h
index d8319ac19..5199bc9e2 100644
--- a/src/broadcast.h
+++ b/src/broadcast.h
@@ -19,12 +19,12 @@ static APY_INLINE bool is_broadcastable(
     const std::vector<std::size_t>& src_shape, const std::vector<std::size_t>& dst_shape
 )
 {
-    // Can not broadcast if either shape is zero-dimensional
+    // Cannot broadcast if either shape is zero-dimensional
     if (src_shape.size() == 0 || dst_shape.size() == 0) {
         return false;
     }
 
-    // Can not broadcast if either shape has a zero dimension
+    // Cannot broadcast if either shape has a zero dimension
     auto is_zero = [](auto n) { return n == 0; };
     for (const auto& shape : { src_shape, dst_shape }) {
         if (std::any_of(std::begin(shape), std::end(shape), is_zero)) {
@@ -32,7 +32,7 @@ static APY_INLINE bool is_broadcastable(
         }
     }
 
-    // Can not broadcast if destination shape has fewer dimensions than source shape
+    // Cannot broadcast if destination shape has fewer dimensions than source shape
     if (src_shape.size() > dst_shape.size()) {
         return false;
     }


Test patch:
diff --git a/lib/test/test_set_item.py b/lib/test/test_set_item.py
new file mode 100644
index 000000000..dfbfea456
--- /dev/null
+++ b/lib/test/test_set_item.py
@@ -0,0 +1,184 @@
+from apytypes import APyFixedArray, APyFloatArray
+from apytypes import APyFixed, APyFloat
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_raise_invalid_item(APyArray, APyScalar):
+    array = APyArray.from_float([[1, 2, 3], [4, 5, 6]], 10, 10)
+    with pytest.raises(
+        ValueError, match=r"APy(Fixed|Float)Array\.__setitem__: supported keys"
+    ):
+        array[1, 2.5] = APyScalar.from_float(3.0, 10, 10)
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_raises_same_type(APyArray, APyScalar):
+    array = APyArray.from_float([[1, 2, 3], [4, 5, 6]], 10, 10)
+    with pytest.raises(
+        ValueError, match=r"APy(Fixed|Float)Array\.__setitem__: `val` has different"
+    ):
+        array[0] = APyScalar.from_float(0, 11, 10)
+
+    with pytest.raises(
+        ValueError, match=r"APy(Fixed|Float)Array\.__setitem__: `val` has different"
+    ):
+        array[0] = APyArray.from_float([0], 11, 10)
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_set_item_integer(APyArray, APyScalar):
+    # ndim == 1
+    array = APyArray.from_float([1, 2, 3, 4, 5, 6], 10, 10)
+
+    array[0] = APyScalar.from_float(10, 10, 10)
+    assert array.is_identical(APyArray.from_float([10, 2, 3, 4, 5, 6], 10, 10))
+    array[5] = APyScalar.from_float(15, 10, 10)
+    assert array.is_identical(APyArray.from_float([10, 2, 3, 4, 5, 15], 10, 10))
+    array[3] = APyScalar.from_float(20, 10, 10)
+    assert array.is_identical(APyArray.from_float([10, 2, 3, 20, 5, 15], 10, 10))
+
+    with pytest.raises(IndexError, match=r"APy(Fixed|Float)Array\.__setitem__"):
+        array[6] = APyScalar.from_float(15, 10, 10)
+
+    # ndim == 2
+    array = APyArray.from_float([[1, 2], [3, 4]], 10, 10)
+
+    array[0] = APyScalar.from_float(10, 10, 10)
+    assert array.is_identical(APyArray.from_float([[10, 10], [3, 4]], 10, 10))
+    array[1] = APyScalar.from_float(12, 10, 10)
+    assert array.is_identical(APyArray.from_float([[10, 10], [12, 12]], 10, 10))
+    array[0] = APyArray.from_float([5, 6], 10, 10)
+    assert array.is_identical(APyArray.from_float([[5, 6], [12, 12]], 10, 10))
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_set_item_slice(APyArray, APyScalar):
+    # ndim == 1
+    array = APyArray.from_float([1, 2, 3, 4, 5, 6], 10, 10)
+
+    array[0:2] = APyScalar.from_float(-3, 10, 10)
+    assert array.is_identical(APyArray.from_float([-3, -3, 3, 4, 5, 6], 10, 10))
+    array[1:4] = APyScalar.from_float(-1, 10, 10)
+    assert array.is_identical(APyArray.from_float([-3, -1, -1, -1, 5, 6], 10, 10))
+    array[::-1] = APyArray.from_float([1, 2, 3, 4, 5, 6], 10, 10)
+    assert array.is_identical(APyArray.from_float([6, 5, 4, 3, 2, 1], 10, 10))
+
+    # ndim == 2
+    array = APyArray.from_float([[1, 2], [3, 4], [5, 6], [7, 8]], 10, 10)
+
+    array[1:3] = APyScalar.from_float(0, 10, 10)
+    assert array.is_identical(
+        APyArray.from_float([[1, 2], [0, 0], [0, 0], [7, 8]], 10, 10)
+    )
+
+    array[::-1] = APyArray.from_float([4, 4], 10, 10)
+    assert array.is_identical(
+        APyArray.from_float([[4, 4], [4, 4], [4, 4], [4, 4]], 10, 10)
+    )
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_set_item_single_ellipsis(APyArray, APyScalar):
+    # ndim == 1
+    array = APyArray.from_float([1, 2, 3, 4, 5, 6], 10, 10)
+
+    array[...] = APyScalar.from_float(123, 10, 10)
+    assert array.is_identical(
+        APyArray.from_float([123, 123, 123, 123, 123, 123], 10, 10)
+    )
+    array[...] = APyArray.from_float([0, 1, 2, 3, 4, 5], 10, 10)
+    assert array.is_identical(APyArray.from_float([0, 1, 2, 3, 4, 5], 10, 10))
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_set_item_tuple(APyArray, APyScalar):
+    np = pytest.importorskip("numpy")
+    np_array = np.array(range(3 * 4 * 5 * 6 * 7)).reshape((3, 4, 5, 6, 7))
+    ap_array = APyArray.from_float(np_array, 25, 25)
+    assert ap_array.shape == np_array.shape
+    assert ap_array[...].is_identical(ap_array)
+
+    ap_array[2, ..., 4] = APyScalar.from_float(10, 25, 25)
+    np_array[2, ..., 4] = 10
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[..., 1, 4, 3] = np.array(range(4))
+    ap_array[..., 1, 4, 3] = APyArray.from_float(range(4), 25, 25)
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[::-1, ..., 2, 1] = np.array(range(3 * 4 * 5)).reshape((3, 4, 5))
+    ap_array[::-1, ..., 2, 1] = APyArray.from_float(
+        np.array(range(3 * 4 * 5)).reshape((3, 4, 5)),
+        25,
+        25,
+    )
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[2, 1, 6:1:-2, 1:6:2] = np.array(range(2 * 3 * 7)).reshape((2, 3, 7))
+    ap_array[2, 1, 6:1:-2, 1:6:2] = APyArray.from_float(
+        np.array(range(2 * 3 * 7)).reshape((2, 3, 7)),
+        25,
+        25,
+    )
+    assert np.all(ap_array.to_numpy() == np_array)
+
+
+def test_set_item_tuple_multi_limb():
+    np = pytest.importorskip("numpy")
+    np_array = np.array(range(3 * 4 * 5 * 6 * 7)).reshape((3, 4, 5, 6, 7))
+    ap_array = APyFixedArray.from_float(np_array, 500, 25)
+    assert ap_array.shape == np_array.shape
+    assert ap_array[...].is_identical(ap_array)
+
+    ap_array[2, ..., 4] = APyFixed.from_float(10, 500, 25)
+    np_array[2, ..., 4] = 10
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[..., 1, 4, 3] = np.array(range(4))
+    ap_array[..., 1, 4, 3] = APyFixedArray.from_float(range(4), 500, 25)
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[::-1, ..., 2, 1] = np.array(range(3 * 4 * 5)).reshape((3, 4, 5))
+    ap_array[::-1, ..., 2, 1] = APyFixedArray.from_float(
+        np.array(range(3 * 4 * 5)).reshape((3, 4, 5)),
+        500,
+        25,
+    )
+    assert np.all(ap_array.to_numpy() == np_array)
+
+    np_array[2, 1, 6:1:-2, 1:6:2] = np.array(range(2 * 3 * 7)).reshape((2, 3, 7))
+    ap_array[2, 1, 6:1:-2, 1:6:2] = APyFixedArray.from_float(
+        np.array(range(2 * 3 * 7)).reshape((2, 3, 7)),
+        500,
+        25,
+    )
+    assert np.all(ap_array.to_numpy() == np_array)
+
+
+@pytest.mark.parametrize(
+    "APyArray, APyScalar", [(APyFixedArray, APyFixed), (APyFloatArray, APyFloat)]
+)
+def test_set_item_not_broadcastable(APyArray, APyScalar):
+    np = pytest.importorskip("numpy")
+    np_array = np.array(range(3 * 4 * 5 * 6 * 7)).reshape((3, 4, 5, 6, 7))
+    ap_array = APyArray.from_float(np_array, 25, 25)
+
+    with pytest.raises(
+        ValueError,
+        match=r"APy(Fixed|Float)Array.__setitem__: `val` shape not broadcastable",
+    ):
+        ap_array[1, 2, 3, 4] = APyArray.from_float([1, 2], 25, 25)