Repository: apytypes/apytypes (language: python)

Issue:
Implement cast context for APyFixed
Implement a dedicated context for casting fixed-point numbers, named `APyFixedCastContext`. It should take a quantization mode and an overflow mode as parameters, potentially a seed as well for future use.


Discussion / hints:

Fix patch:
diff --git a/docs/api/context.rst b/docs/api/context.rst
index 29024ae24..1631d210e 100644
--- a/docs/api/context.rst
+++ b/docs/api/context.rst
@@ -1,6 +1,10 @@
 Context handling
 ================
 
+.. autoclass:: apytypes.APyFixedCastContext
+
+   .. automethod:: __init__
+
 .. autoclass:: apytypes.APyFloatQuantizationContext
 
    .. automethod:: __init__
diff --git a/lib/apytypes/__init__.py b/lib/apytypes/__init__.py
index 3980091da..3913ebd79 100644
--- a/lib/apytypes/__init__.py
+++ b/lib/apytypes/__init__.py
@@ -4,6 +4,7 @@
     APyFixed,
     APyFixedArray,
     APyFixedAccumulatorContext,
+    APyFixedCastContext,
     APyFloat,
     APyFloatArray,
     APyFloatAccumulatorContext,
@@ -23,6 +24,7 @@
     "APyFixed",
     "APyFixedArray",
     "APyFixedAccumulatorContext",
+    "APyFixedCastContext",
     "APyFloat",
     "APyFloatArray",
     "APyFloatAccumulatorContext",
diff --git a/lib/apytypes/__init__.pyi b/lib/apytypes/__init__.pyi
index 408c82329..e75cf1d7a 100644
--- a/lib/apytypes/__init__.pyi
+++ b/lib/apytypes/__init__.pyi
@@ -2,6 +2,7 @@ from __future__ import annotations
 from apytypes._apytypes import APyFixed
 from apytypes._apytypes import APyFixedArray
 from apytypes._apytypes import APyFixedAccumulatorContext
+from apytypes._apytypes import APyFixedCastContext
 from apytypes._apytypes import APyFloat
 from apytypes._apytypes import APyFloatArray
 from apytypes._apytypes import APyFloatAccumulatorContext
@@ -19,6 +20,7 @@ __all__: list = [
     "APyFixed",
     "APyFixedArray",
     "APyFixedAccumulatorContext",
+    "APyFixedCastContext",
     "APyFloat",
     "APyFloatArray",
     "APyFloatAccumulatorContext",
diff --git a/lib/apytypes/_apytypes.pyi b/lib/apytypes/_apytypes.pyi
index a02dac710..0a3aed5c2 100644
--- a/lib/apytypes/_apytypes.pyi
+++ b/lib/apytypes/_apytypes.pyi
@@ -8,6 +8,7 @@ __all__ = [
     "APyFixedArray",
     "APyFixedArrayIterator",
     "APyFixedAccumulatorContext",
+    "APyFixedCastContext",
     "APyFloat",
     "APyFloatArray",
     "APyFloatArrayIterator",
@@ -177,8 +178,8 @@ class APyFixed:
         self,
         bits: int | None = None,
         int_bits: int | None = None,
-        quantization: QuantizationMode = QuantizationMode.TRN,
-        overflow: OverflowMode = OverflowMode.WRAP,
+        quantization: QuantizationMode = None,
+        overflow: OverflowMode = None,
         frac_bits: int | None = None,
     ) -> APyFixed:
         """
@@ -479,8 +480,8 @@ class APyFixedArray:
         self,
         bits: int | None = None,
         int_bits: int | None = None,
-        quantization: QuantizationMode = QuantizationMode.TRN,
-        overflow: OverflowMode = OverflowMode.WRAP,
+        quantization: QuantizationMode = None,
+        overflow: OverflowMode = None,
         frac_bits: int | None = None,
     ) -> APyFixedArray:
         """
@@ -1247,6 +1248,21 @@ class APyFixedAccumulatorContext(ContextManager):
         overflow: OverflowMode | None = None,
     ) -> None: ...
 
+class APyFixedCastContext(ContextManager):
+    def __enter__(self: ContextManager) -> None: ...
+    def __exit__(
+        self: ContextManager,
+        arg0: type | None,
+        arg1: typing.Any | None,
+        arg2: typing.Any | None,
+    ) -> None: ...
+    def __init__(
+        self,
+        *,
+        quantization: QuantizationMode | None = None,
+        overflow: OverflowMode | None = None,
+    ) -> None: ...
+
 class ContextManager:
     pass
 
diff --git a/src/apyfixed.cc b/src/apyfixed.cc
index 2e181d6e8..42a593955 100644
--- a/src/apyfixed.cc
+++ b/src/apyfixed.cc
@@ -26,6 +26,7 @@ namespace nb = nanobind;
 
 #include "apyfixed.h"
 #include "apyfloat.h"
+#include "apytypes_common.h"
 #include "apytypes_util.h"
 #include "ieee754.h"
 #include "python_util.h"
@@ -926,8 +927,8 @@ APyFixed APyFixed::from_string(
 APyFixed APyFixed::cast(
     std::optional<int> bits,
     std::optional<int> int_bits,
-    QuantizationMode quantization,
-    OverflowMode overflow,
+    std::optional<QuantizationMode> quantization,
+    std::optional<OverflowMode> overflow,
     std::optional<int> frac_bits
 ) const
 {
@@ -935,6 +936,10 @@ APyFixed APyFixed::cast(
     int new_bits = bits_from_optional(bits, int_bits, frac_bits);
     int new_int_bits = int_bits.has_value() ? *int_bits : *bits - *frac_bits;
 
+    const APyFixedCastOption cast_option = get_fixed_cast_mode();
+    const auto quantization_mode = quantization.value_or(cast_option.quantization);
+    const auto overflow_mode = overflow.value_or(cast_option.overflow);
+
     // Result that temporarily can hold all the necessary bits
     APyFixed result(std::max(new_bits, _bits), new_int_bits);
     _cast(
@@ -942,8 +947,8 @@ APyFixed APyFixed::cast(
         result._data.end(),   // output sentinel
         new_bits,
         new_int_bits,
-        quantization,
-        overflow
+        quantization_mode,
+        overflow_mode
     );
 
     result._bits = new_bits;
diff --git a/src/apyfixed.h b/src/apyfixed.h
index 48770d3cc..997089498 100644
--- a/src/apyfixed.h
+++ b/src/apyfixed.h
@@ -270,8 +270,8 @@ class APyFixed {
     APyFixed cast(
         std::optional<int> bits = std::nullopt,
         std::optional<int> int_bits = std::nullopt,
-        QuantizationMode quantization = QuantizationMode::TRN,
-        OverflowMode overflow = OverflowMode::WRAP,
+        std::optional<QuantizationMode> quantization = std::nullopt,
+        std::optional<OverflowMode> overflow = std::nullopt,
         std::optional<int> frac_bits = std::nullopt
     ) const;
 
diff --git a/src/apyfixed_wrapper.cc b/src/apyfixed_wrapper.cc
index 93a8605b6..4b0220b60 100644
--- a/src/apyfixed_wrapper.cc
+++ b/src/apyfixed_wrapper.cc
@@ -167,8 +167,8 @@ void bind_fixed(nb::module_& m)
             &APyFixed::cast,
             nb::arg("bits") = nb::none(),
             nb::arg("int_bits") = nb::none(),
-            nb::arg("quantization") = QuantizationMode::TRN,
-            nb::arg("overflow") = OverflowMode::WRAP,
+            nb::arg("quantization") = nb::none(),
+            nb::arg("overflow") = nb::none(),
             nb::arg("frac_bits") = nb::none(),
             R"pbdoc(
             Change format of the fixed-point number.
diff --git a/src/apyfixedarray.cc b/src/apyfixedarray.cc
index bae0f4e18..4cb04598d 100644
--- a/src/apyfixedarray.cc
+++ b/src/apyfixedarray.cc
@@ -919,8 +919,8 @@ bool APyFixedArray::is_identical(const APyFixedArray& other) const
 APyFixedArray APyFixedArray::cast(
     std::optional<int> bits,
     std::optional<int> int_bits,
-    QuantizationMode quantization,
-    OverflowMode overflow,
+    std::optional<QuantizationMode> quantization,
+    std::optional<OverflowMode> overflow,
     std::optional<int> frac_bits
 ) const
 {
@@ -928,6 +928,10 @@ APyFixedArray APyFixedArray::cast(
     int new_bits = bits_from_optional(bits, int_bits, frac_bits);
     int new_int_bits = int_bits.has_value() ? *int_bits : *bits - *frac_bits;
 
+    const APyFixedCastOption cast_option = get_fixed_cast_mode();
+    const auto quantization_mode = quantization.value_or(cast_option.quantization);
+    const auto overflow_mode = overflow.value_or(cast_option.overflow);
+
     // The new result array (`bit_specifier_sanitize()` called in constructor)
     std::size_t result_limbs = bits_to_limbs(new_bits);
     std::size_t pad_limbs = bits_to_limbs(std::max(new_bits, _bits)) - result_limbs;
@@ -947,8 +951,8 @@ APyFixedArray APyFixedArray::cast(
         caster,
         new_bits,
         new_int_bits,
-        quantization,
-        overflow
+        quantization_mode,
+        overflow_mode
     );
 
     result._bits = new_bits;
diff --git a/src/apyfixedarray.h b/src/apyfixedarray.h
index 10f8b1473..fc1f760a4 100644
--- a/src/apyfixedarray.h
+++ b/src/apyfixedarray.h
@@ -159,8 +159,8 @@ class APyFixedArray : public APyBuffer<mp_limb_t> {
     APyFixedArray cast(
         std::optional<int> bits = std::nullopt,
         std::optional<int> int_bits = std::nullopt,
-        QuantizationMode quantization = QuantizationMode::TRN,
-        OverflowMode overflow = OverflowMode::WRAP,
+        std::optional<QuantizationMode> quantization = QuantizationMode::TRN,
+        std::optional<OverflowMode> overflow = OverflowMode::WRAP,
         std::optional<int> frac_bits = std::nullopt
     ) const;
 
diff --git a/src/apyfixedarray_wrapper.cc b/src/apyfixedarray_wrapper.cc
index 4e3bec209..8ebd0e2c1 100644
--- a/src/apyfixedarray_wrapper.cc
+++ b/src/apyfixedarray_wrapper.cc
@@ -307,8 +307,8 @@ void bind_fixed_array(nb::module_& m)
             &APyFixedArray::cast,
             nb::arg("bits") = nb::none(),
             nb::arg("int_bits") = nb::none(),
-            nb::arg("quantization") = QuantizationMode::TRN,
-            nb::arg("overflow") = OverflowMode::WRAP,
+            nb::arg("quantization") = nb::none(),
+            nb::arg("overflow") = nb::none(),
             nb::arg("frac_bits") = nb::none(),
             R"pbdoc(
             Change format of the fixed-point array.
diff --git a/src/apytypes_common.cc b/src/apytypes_common.cc
index eaa25bbda..20b99037f 100644
--- a/src/apytypes_common.cc
+++ b/src/apytypes_common.cc
@@ -74,6 +74,42 @@ std::uint64_t get_float_quantization_seed() { return quantization_seed; }
 
 std::uint64_t random_number_float() { return gen64(); }
 
+/* ********************************************************************************** *
+ * *                      Cast context for APyFixed                                 * *
+ * ********************************************************************************** */
+
+// Global accumulator option (default value: std::nullopt)
+static APyFixedCastOption global_cast_option_fixed
+    = { QuantizationMode::TRN, OverflowMode::WRAP };
+
+APyFixedCastContext::APyFixedCastContext(
+    std::optional<QuantizationMode> quantization, std::optional<OverflowMode> overflow
+)
+{
+    if (!quantization.has_value() && !overflow.has_value()) {
+        throw nb::value_error(
+            "Either quantization mode or overflow mode must be specified."
+        );
+    }
+
+    // Store the previous accumulator mode
+    previous_mode = global_cast_option_fixed;
+
+    // Extract the input
+    APyFixedCastOption new_mode;
+
+    new_mode.quantization = quantization.value_or(previous_mode.quantization);
+    new_mode.overflow = overflow.value_or(previous_mode.overflow);
+
+    // Set the current mode
+    current_mode = new_mode;
+}
+
+void APyFixedCastContext::enter_context() { global_cast_option_fixed = current_mode; }
+void APyFixedCastContext::exit_context() { global_cast_option_fixed = previous_mode; }
+
+APyFixedCastOption get_fixed_cast_mode() { return global_cast_option_fixed; }
+
 /* ********************************************************************************** *
  * *                      Accumulator context for APyFixedArray                     * *
  * ********************************************************************************** */
diff --git a/src/apytypes_common.h b/src/apytypes_common.h
index 0b9d8abb4..d22d974e6 100644
--- a/src/apytypes_common.h
+++ b/src/apytypes_common.h
@@ -108,7 +108,7 @@ class ContextManager {
 };
 
 /* ********************************************************************************** *
- * *                          Quantization context for APyFloat * *
+ * *                          Quantization context for APyFloat                     * *
  * ********************************************************************************** */
 
 /*
@@ -174,6 +174,31 @@ struct APyFloatData {
     }
 };
 
+/* ********************************************************************************** *
+ * *                          Cast context for APyFixed                             * *
+ * ********************************************************************************** */
+
+struct APyFixedCastOption {
+    QuantizationMode quantization;
+    OverflowMode overflow;
+};
+
+class APyFixedCastContext : public ContextManager {
+public:
+    APyFixedCastContext(
+        std::optional<QuantizationMode> quantization = std::nullopt,
+        std::optional<OverflowMode> overflow = std::nullopt
+    );
+    void enter_context() override;
+    void exit_context() override;
+
+private:
+    APyFixedCastOption current_mode, previous_mode;
+};
+
+//! Retrieve the global cast mode for APyFixed
+APyFixedCastOption get_fixed_cast_mode();
+
 /* ********************************************************************************** *
  * *                          Accumulator context for APyFixed                      * *
  * ********************************************************************************** */
diff --git a/src/apytypes_context_wrapper.cc b/src/apytypes_context_wrapper.cc
index dae3da43f..e7a963a40 100644
--- a/src/apytypes_context_wrapper.cc
+++ b/src/apytypes_context_wrapper.cc
@@ -45,6 +45,26 @@ void bind_quantization_context(nb::module_& m)
         );
 }
 
+void bind_cast_context(nb::module_& m)
+{
+    nb::class_<APyFixedCastContext, ContextManager>(m, "APyFixedCastContext")
+        .def(
+            nb::init<std::optional<QuantizationMode>, std::optional<OverflowMode>>(),
+            // nb::kw_only() is added in NanoBind v2.0.0
+            // nb::kw_only(), // All parameters are keyword only
+            nb::arg("quantization") = nb::none(),
+            nb::arg("overflow") = nb::none()
+        )
+        .def("__enter__", &context_enter_handler)
+        .def(
+            "__exit__",
+            &context_exit_handler,
+            nb::arg("exc_type") = nb::none(),
+            nb::arg("exc_value") = nb::none(),
+            nb::arg("traceback") = nb::none()
+        );
+}
+
 void bind_accumulator_context(nb::module_& m)
 {
     nb::class_<APyFixedAccumulatorContext, ContextManager>(
diff --git a/src/apytypes_wrapper.cc b/src/apytypes_wrapper.cc
index 344cb0d89..b218a5c67 100644
--- a/src/apytypes_wrapper.cc
+++ b/src/apytypes_wrapper.cc
@@ -5,6 +5,7 @@
 namespace nb = nanobind;
 
 void bind_accumulator_context(nb::module_& m);
+void bind_cast_context(nb::module_& m);
 void bind_common(nb::module_& m);
 void bind_context_manager(nb::module_& m);
 void bind_fixed(nb::module_& m);
@@ -23,6 +24,7 @@ NB_MODULE(_apytypes, m)
     bind_context_manager(m);
     bind_quantization_context(m);
     bind_accumulator_context(m);
+    bind_cast_context(m);
 }
 
 void bind_common(nb::module_& m)


Test patch:
diff --git a/lib/test/test_contextmanagers.py b/lib/test/test_contextmanagers.py
index 766dd67c2..cad0aa046 100644
--- a/lib/test/test_contextmanagers.py
+++ b/lib/test/test_contextmanagers.py
@@ -1,10 +1,5 @@
 import apytypes
-from apytypes import (
-    APyFloatQuantizationContext,
-    QuantizationMode,
-    APyFloatAccumulatorContext,
-    APyFixedAccumulatorContext,
-)
+from apytypes import *
 import pytest
 
 
@@ -19,6 +14,199 @@ def test_context_kw_only():
             pass
 
 
+class TestCastContext:
+    """
+    This test class doesn't test if cast itself works,
+    just that the context manager acts correctly.
+    """
+
+    def test_raises(self):
+        with pytest.raises(ValueError, match=".*must be specified"):
+            with APyFixedCastContext():
+                pass
+
+        with APyFixedCastContext(quantization=QuantizationMode.TRN):  # should not raise
+            pass
+
+        with APyFixedCastContext(overflow=OverflowMode.SAT):  # should not raise
+            pass
+
+        with APyFixedCastContext(
+            quantization=QuantizationMode.TRN, overflow=OverflowMode.SAT
+        ):  # should not raise
+            pass
+
+    def test_with_cast_context(self):
+        a = APyFixed.from_float(1.25, int_bits=3, frac_bits=2)
+        # TRN by default
+        assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+            APyFixed(2, bits=4, int_bits=3)
+        )
+
+        with APyFixedCastContext(quantization=QuantizationMode.TRN_INF):
+            assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+                APyFixed(3, bits=4, int_bits=3)
+            )
+
+        # TRN by default again
+        assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+            APyFixed(2, bits=4, int_bits=3)
+        )
+
+    def test_cast_context_override(self):
+        a = APyFixed.from_float(1.25, int_bits=3, frac_bits=2)
+
+        with APyFixedCastContext(quantization=QuantizationMode.TRN_INF):
+            assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+                APyFixed(3, bits=4, int_bits=3)
+            )
+
+            # Use TRN explicitly inside context
+            assert (
+                _ := a.cast(int_bits=3, frac_bits=1, quantization=QuantizationMode.TRN)
+            ).is_identical(APyFixed(2, bits=4, int_bits=3))
+
+    def test_cast_context_nested(self):
+        a = APyFixed.from_float(1.25, int_bits=3, frac_bits=2)
+        # TRN by default
+        assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+            APyFixed(2, bits=4, int_bits=3)
+        )
+
+        with APyFixedCastContext(quantization=QuantizationMode.TRN_INF):
+            assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+                APyFixed(3, bits=4, int_bits=3)
+            )
+
+            with APyFixedCastContext(quantization=QuantizationMode.TRN):
+                assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+                    APyFixed(2, bits=4, int_bits=3)
+                )
+
+            assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+                APyFixed(3, bits=4, int_bits=3)
+            )
+
+        assert (_ := a.cast(int_bits=3, frac_bits=1)).is_identical(
+            APyFixed(2, bits=4, int_bits=3)
+        )
+
+    def test_cast_context_overflow(self):
+        """
+        If the previous tests pass for the quantization mode, then overflow should also work.
+        These are therefore more like sanity tests.
+        """
+        a = APyFixed.from_float(4, int_bits=4, frac_bits=0)
+        # WRAP by default
+        assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+            APyFixed(4, bits=3, int_bits=3)
+        )
+
+        # Set only overflow mode
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixed(3, bits=3, int_bits=3)
+            )
+
+        # Override overflow mode
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (
+                _ := a.cast(int_bits=3, frac_bits=0, overflow=OverflowMode.WRAP)
+            ).is_identical(APyFixed(4, bits=3, int_bits=3))
+
+        # Nested contexts
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixed(3, bits=3, int_bits=3)
+            )
+
+            with APyFixedCastContext(overflow=OverflowMode.WRAP):
+                assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                    APyFixed(4, bits=3, int_bits=3)
+                )
+
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixed(3, bits=3, int_bits=3)
+            )
+
+        # Set both quantization mode and overflow mode
+        a = APyFixed.from_float(3.5, int_bits=3, frac_bits=1)
+        quantization_mode = QuantizationMode.TRN_INF
+        overflow_mode = OverflowMode.SAT
+        assert (
+            _ := a.cast(
+                int_bits=3,
+                frac_bits=0,
+                quantization=quantization_mode,
+                overflow=overflow_mode,
+            )
+        ).is_identical(APyFixed(3, int_bits=3, frac_bits=0))
+
+        with APyFixedCastContext(
+            quantization=quantization_mode, overflow=overflow_mode
+        ):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixed(3, int_bits=3, frac_bits=0)
+            )
+
+    def test_cast_context_array(self):
+        """
+        Basic sanity tests, if the scalar case works then it should work for arrays as well.
+        """
+        a = APyFixedArray.from_float([4], int_bits=4, frac_bits=0)
+        # WRAP by default
+        assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+            APyFixedArray([4], bits=3, int_bits=3)
+        )
+
+        # Set only overflow mode
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixedArray([3], bits=3, int_bits=3)
+            )
+
+        # Override overflow mode
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (
+                _ := a.cast(int_bits=3, frac_bits=0, overflow=OverflowMode.WRAP)
+            ).is_identical(APyFixedArray([4], bits=3, int_bits=3))
+
+        # Nested contexts
+        with APyFixedCastContext(overflow=OverflowMode.SAT):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixedArray([3], bits=3, int_bits=3)
+            )
+
+            with APyFixedCastContext(overflow=OverflowMode.WRAP):
+                assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                    APyFixedArray([4], bits=3, int_bits=3)
+                )
+
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixedArray([3], bits=3, int_bits=3)
+            )
+
+        # Set both quantization mode and overflow mode
+        a = APyFixedArray.from_float([3.5], int_bits=3, frac_bits=1)
+        quantization_mode = QuantizationMode.TRN_INF
+        overflow_mode = OverflowMode.SAT
+        assert (
+            _ := a.cast(
+                int_bits=3,
+                frac_bits=0,
+                quantization=quantization_mode,
+                overflow=overflow_mode,
+            )
+        ).is_identical(APyFixedArray([3], int_bits=3, frac_bits=0))
+
+        with APyFixedCastContext(
+            quantization=quantization_mode, overflow=overflow_mode
+        ):
+            assert (_ := a.cast(int_bits=3, frac_bits=0)).is_identical(
+                APyFixedArray([3], int_bits=3, frac_bits=0)
+            )
+
+
 class TestAccumulatorContext:
     """
     This test class doesn't test if the accumulation itself works,