slbotbm commented on code in PR #3613:
URL: https://github.com/apache/iggy/pull/3613#discussion_r3677316623
##########
foreign/python/src/receive_message.rs:
##########
@@ -72,6 +81,19 @@ impl ReceiveMessage {
pub fn partition_id(&self) -> u32 {
self.partition_id
}
+
+ /// Retrieves user headers attached to the received message.
+ #[gen_stub(override_return_type(type_repr = "UserHeaders | None"))]
+ pub fn user_headers<'a>(&self, py: Python<'a>) ->
PyResult<Option<Bound<'a, UserHeaders>>> {
+ let Some(headers) = self
+ .inner
+ .user_headers_map()
+ .map_err(|e| PyValueError::new_err(e.to_string()))?
+ else {
+ return Ok(None);
+ };
+ rust_user_headers_to_py(py, headers).map(Some)
+ }
Review Comment:
The linked issue notes that structurally malformed wire headers are logged
and
returned as `None`, making them indistinguishable from absent headers.
Document the `None` behavior and distinguish it from known semantic decode
errors, which become `ValueError`.
##########
foreign/python/src/user_headers.rs:
##########
@@ -0,0 +1,897 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::BTreeMap;
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+
+use iggy::prelude::{
+ HeaderField, HeaderKey as RustHeaderKey, HeaderKind, HeaderValue as
RustHeaderValue,
+};
+use pyo3::PyClass;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::pyclass::CompareOp;
+use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyString};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum,
gen_stub_pymethods};
+
+type RustUserHeaders = BTreeMap<RustHeaderKey, RustHeaderValue>;
+
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+/// Typed key for an Iggy user header.
+///
+/// Use these constructors when the header key must preserve an explicit
+/// wire type instead of using the common string-key dictionary form.
+pub enum HeaderKey {
+ /// Raw bytes key. The byte length must be 1..=255.
+ Raw { value: Py<PyBytes> },
+ /// UTF-8 string key. The encoded byte length must be 1..=255.
+ String { value: String },
+ /// Boolean key.
+ Bool { value: bool },
+ /// Signed 8-bit integer key.
+ Int8 { value: i8 },
+ /// Signed 16-bit integer key.
+ Int16 { value: i16 },
+ /// Signed 32-bit integer key.
+ Int32 { value: i32 },
+ /// Signed 64-bit integer key.
+ Int64 { value: i64 },
+ /// Signed 128-bit integer key.
+ Int128 { value: i128 },
+ /// Unsigned 8-bit integer key.
+ UnsignedInt8 { value: u8 },
+ /// Unsigned 16-bit integer key.
+ UnsignedInt16 { value: u16 },
+ /// Unsigned 32-bit integer key.
+ UnsignedInt32 { value: u32 },
+ /// Unsigned 64-bit integer key.
+ UnsignedInt64 { value: u64 },
+ /// Unsigned 128-bit integer key.
+ UnsignedInt128 { value: u128 },
+ /// 32-bit floating point key.
+ Float32 { value: f32 },
+ /// 64-bit floating point key.
+ Float64 { value: f64 },
+}
+
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+/// Typed value for an Iggy user header.
+///
+/// Use these constructors when the header value must preserve an explicit
+/// wire type instead of using the common Python scalar dictionary form.
+pub enum HeaderValue {
+ /// Raw bytes value. The byte length must be 1..=255.
+ Raw { value: Py<PyBytes> },
+ /// UTF-8 string value. The encoded byte length must be 1..=255.
+ String { value: String },
+ /// Boolean value.
+ Bool { value: bool },
+ /// Signed 8-bit integer value.
+ Int8 { value: i8 },
+ /// Signed 16-bit integer value.
+ Int16 { value: i16 },
+ /// Signed 32-bit integer value.
+ Int32 { value: i32 },
+ /// Signed 64-bit integer value.
+ Int64 { value: i64 },
+ /// Signed 128-bit integer value.
+ Int128 { value: i128 },
+ /// Unsigned 8-bit integer value.
+ UnsignedInt8 { value: u8 },
+ /// Unsigned 16-bit integer value.
+ UnsignedInt16 { value: u16 },
+ /// Unsigned 32-bit integer value.
+ UnsignedInt32 { value: u32 },
+ /// Unsigned 64-bit integer value.
+ UnsignedInt64 { value: u64 },
+ /// Unsigned 128-bit integer value.
+ UnsignedInt128 { value: u128 },
+ /// 32-bit floating point value.
+ Float32 { value: f32 },
+ /// 64-bit floating point value.
+ Float64 { value: f64 },
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl HeaderKey {
+ pub fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
+ Ok(py_hash(header_identity_hash(self.identity(py)?)))
+ }
+
+ pub fn __richcmp__(
+ &self,
+ py: Python<'_>,
+ other: &Bound<'_, PyAny>,
+ op: CompareOp,
+ ) -> PyResult<Py<PyAny>> {
+ header_key_richcmp(py, self.identity(py)?, other, op)
+ }
+
+ pub fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+ self.repr(py)
+ }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl HeaderValue {
+ pub fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
+ Ok(py_hash(header_identity_hash(self.identity(py)?)))
+ }
+
+ pub fn __richcmp__(
+ &self,
+ py: Python<'_>,
+ other: &Bound<'_, PyAny>,
+ op: CompareOp,
+ ) -> PyResult<Py<PyAny>> {
+ header_value_richcmp(py, self.identity(py)?, other, op)
+ }
+
+ pub fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+ self.repr(py)
+ }
+}
+
+/// Borrows a typed Python [`HeaderKey`] together with the GIL token so it can
+/// be converted into its Rust representation (the `Raw` variant needs the
token
+/// to read its `PyBytes` payload).
+struct PyHeaderKeyRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value HeaderKey,
+}
+
+impl TryFrom<PyHeaderKeyRef<'_, '_>> for RustHeaderKey {
+ type Error = PyErr;
+
+ fn try_from(key: PyHeaderKeyRef<'_, '_>) -> PyResult<Self> {
+ let PyHeaderKeyRef { py, value } = key;
+ match value {
+ HeaderKey::Raw { value } => {
+
RustHeaderKey::try_from(value.extract::<Vec<u8>>(py)?).map_err(to_value_error)
+ }
+ HeaderKey::String { value } => {
+ RustHeaderKey::try_from(value.as_str()).map_err(to_value_error)
+ }
+ HeaderKey::Bool { value } => Ok((*value).into()),
+ HeaderKey::Int8 { value } => Ok((*value).into()),
+ HeaderKey::Int16 { value } => Ok((*value).into()),
+ HeaderKey::Int32 { value } => Ok((*value).into()),
+ HeaderKey::Int64 { value } => Ok((*value).into()),
+ HeaderKey::Int128 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt8 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt16 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt32 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt64 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt128 { value } => Ok((*value).into()),
+ HeaderKey::Float32 { value } => checked_float32(*value),
+ HeaderKey::Float64 { value } => Ok((*value).into()),
+ }
+ }
+}
+
+/// Borrows a Rust [`RustHeaderKey`] together with the GIL token so it can be
+/// converted into a typed Python [`HeaderKey`] (the `Raw` variant needs the
+/// token to allocate its `PyBytes` payload).
+struct RustHeaderKeyRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value RustHeaderKey,
+}
+
+impl<'py> TryFrom<RustHeaderKeyRef<'py, '_>> for Bound<'py, HeaderKey> {
+ type Error = PyErr;
+
+ fn try_from(key: RustHeaderKeyRef<'py, '_>) -> PyResult<Self> {
+ let RustHeaderKeyRef { py, value } = key;
+ let key = match value.kind() {
+ HeaderKind::Raw => HeaderKey::Raw {
+ value: PyBytes::new(py,
value.as_raw().map_err(to_value_error)?).unbind(),
+ },
+ HeaderKind::String => HeaderKey::String {
+ value: value.as_str().map_err(to_value_error)?.to_string(),
+ },
+ HeaderKind::Bool => HeaderKey::Bool {
+ value: value.as_bool().map_err(to_value_error)?,
+ },
+ HeaderKind::Int8 => HeaderKey::Int8 {
+ value: value.as_int8().map_err(to_value_error)?,
+ },
+ HeaderKind::Int16 => HeaderKey::Int16 {
+ value: value.as_int16().map_err(to_value_error)?,
+ },
+ HeaderKind::Int32 => HeaderKey::Int32 {
+ value: value.as_int32().map_err(to_value_error)?,
+ },
+ HeaderKind::Int64 => HeaderKey::Int64 {
+ value: value.as_int64().map_err(to_value_error)?,
+ },
+ HeaderKind::Int128 => HeaderKey::Int128 {
+ value: value.as_int128().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint8 => HeaderKey::UnsignedInt8 {
+ value: value.as_uint8().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint16 => HeaderKey::UnsignedInt16 {
+ value: value.as_uint16().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint32 => HeaderKey::UnsignedInt32 {
+ value: value.as_uint32().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint64 => HeaderKey::UnsignedInt64 {
+ value: value.as_uint64().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint128 => HeaderKey::UnsignedInt128 {
+ value: value.as_uint128().map_err(to_value_error)?,
+ },
+ HeaderKind::Float32 => HeaderKey::Float32 {
+ value: value.as_float32().map_err(to_value_error)?,
+ },
+ HeaderKind::Float64 => HeaderKey::Float64 {
+ value: value.as_float64().map_err(to_value_error)?,
+ },
+ };
+ key.into_pyobject(py)
+ }
+}
+
+trait ToPlainScalar {
+ fn to_plain_scalar<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py,
PyAny>>;
+}
+
+macro_rules! impl_to_plain_scalar {
+ ($ty:ty) => {
+ impl ToPlainScalar for $ty {
+ fn to_plain_scalar<'py>(&self, py: Python<'py>) ->
PyResult<Bound<'py, PyAny>> {
+ match self {
+ Self::Raw { value } =>
Ok(value.bind(py).clone().into_any()),
+ Self::String { value } =>
Ok(value.clone().into_pyobject(py)?.into_any()),
+ Self::Bool { value } =>
Ok((*value).into_pyobject(py)?.to_owned().into_any()),
+ Self::Int8 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int16 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int32 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int128 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt8 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt16 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt32 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt128 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Float32 { value } =>
Ok(f64::from(*value).into_pyobject(py)?.into_any()),
+ Self::Float64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ }
+ }
+ }
+ };
+}
+
+impl_to_plain_scalar!(HeaderKey);
+impl_to_plain_scalar!(HeaderValue);
+
+/// Borrows a typed header ([`HeaderKey`] or [`HeaderValue`]) together with the
+/// GIL token for direct conversion to a plain Python scalar.
+struct HeaderToPlainRef<'py, 'value, T: ToPlainScalar> {
+ py: Python<'py>,
+ value: &'value T,
+}
+
+impl<'py, T: ToPlainScalar> TryFrom<HeaderToPlainRef<'py, '_, T>> for
Bound<'py, PyAny> {
+ type Error = PyErr;
+
+ fn try_from(inner: HeaderToPlainRef<'py, '_, T>) -> PyResult<Self> {
+ let HeaderToPlainRef { py, value } = inner;
+ value.to_plain_scalar(py)
+ }
+}
+
+impl HeaderKey {
+ fn identity(&self, py: Python<'_>) -> PyResult<HeaderIdentity> {
+ match self {
+ HeaderKey::Raw { value } =>
Ok(HeaderIdentity::raw(value.extract::<Vec<u8>>(py)?)),
+ HeaderKey::String { value } =>
Ok(HeaderIdentity::string(value.as_bytes().to_vec())),
+ HeaderKey::Bool { value } => Ok(HeaderIdentity::bool(*value)),
+ HeaderKey::Int8 { value } => Ok(HeaderIdentity::int8(*value)),
+ HeaderKey::Int16 { value } => Ok(HeaderIdentity::int16(*value)),
+ HeaderKey::Int32 { value } => Ok(HeaderIdentity::int32(*value)),
+ HeaderKey::Int64 { value } => Ok(HeaderIdentity::int64(*value)),
+ HeaderKey::Int128 { value } => Ok(HeaderIdentity::int128(*value)),
+ HeaderKey::UnsignedInt8 { value } =>
Ok(HeaderIdentity::uint8(*value)),
+ HeaderKey::UnsignedInt16 { value } =>
Ok(HeaderIdentity::uint16(*value)),
+ HeaderKey::UnsignedInt32 { value } =>
Ok(HeaderIdentity::uint32(*value)),
+ HeaderKey::UnsignedInt64 { value } =>
Ok(HeaderIdentity::uint64(*value)),
+ HeaderKey::UnsignedInt128 { value } =>
Ok(HeaderIdentity::uint128(*value)),
+ HeaderKey::Float32 { value } =>
Ok(HeaderIdentity::float32(*value)),
+ HeaderKey::Float64 { value } =>
Ok(HeaderIdentity::float64(*value)),
+ }
+ }
+
+ fn repr(&self, py: Python<'_>) -> PyResult<String> {
+ match self {
+ HeaderKey::Raw { value } => Ok(format!(
+ "HeaderKey.Raw({})",
+ value.bind(py).repr()?.extract::<String>()?
+ )),
+ HeaderKey::String { value } =>
Ok(format!("HeaderKey.String({value:?})")),
+ HeaderKey::Bool { value } =>
Ok(format!("HeaderKey.Bool({value})")),
+ HeaderKey::Int8 { value } =>
Ok(format!("HeaderKey.Int8({value})")),
+ HeaderKey::Int16 { value } =>
Ok(format!("HeaderKey.Int16({value})")),
+ HeaderKey::Int32 { value } =>
Ok(format!("HeaderKey.Int32({value})")),
+ HeaderKey::Int64 { value } =>
Ok(format!("HeaderKey.Int64({value})")),
+ HeaderKey::Int128 { value } =>
Ok(format!("HeaderKey.Int128({value})")),
+ HeaderKey::UnsignedInt8 { value } =>
Ok(format!("HeaderKey.UnsignedInt8({value})")),
+ HeaderKey::UnsignedInt16 { value } =>
Ok(format!("HeaderKey.UnsignedInt16({value})")),
+ HeaderKey::UnsignedInt32 { value } =>
Ok(format!("HeaderKey.UnsignedInt32({value})")),
+ HeaderKey::UnsignedInt64 { value } =>
Ok(format!("HeaderKey.UnsignedInt64({value})")),
+ HeaderKey::UnsignedInt128 { value } =>
Ok(format!("HeaderKey.UnsignedInt128({value})")),
+ HeaderKey::Float32 { value } =>
Ok(format!("HeaderKey.Float32({value:?})")),
+ HeaderKey::Float64 { value } =>
Ok(format!("HeaderKey.Float64({value:?})")),
+ }
+ }
+}
+
+/// Borrows a typed Python [`HeaderValue`] together with the GIL token so it
can
+/// be converted into its Rust representation (the `Raw` variant needs the
token
+/// to read its `PyBytes` payload).
+struct PyHeaderValueRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value HeaderValue,
+}
+
+impl TryFrom<PyHeaderValueRef<'_, '_>> for RustHeaderValue {
+ type Error = PyErr;
+
+ fn try_from(value: PyHeaderValueRef<'_, '_>) -> PyResult<Self> {
+ let PyHeaderValueRef { py, value } = value;
+ match value {
+ HeaderValue::Raw { value } => {
+
RustHeaderValue::try_from(value.extract::<Vec<u8>>(py)?).map_err(to_value_error)
+ }
+ HeaderValue::String { value } => {
+
RustHeaderValue::try_from(value.as_str()).map_err(to_value_error)
+ }
+ HeaderValue::Bool { value } => Ok((*value).into()),
+ HeaderValue::Int8 { value } => Ok((*value).into()),
+ HeaderValue::Int16 { value } => Ok((*value).into()),
+ HeaderValue::Int32 { value } => Ok((*value).into()),
+ HeaderValue::Int64 { value } => Ok((*value).into()),
+ HeaderValue::Int128 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt8 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt16 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt32 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt64 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt128 { value } => Ok((*value).into()),
+ HeaderValue::Float32 { value } => checked_float32(*value),
+ HeaderValue::Float64 { value } => Ok((*value).into()),
+ }
+ }
+}
+
+/// Borrows a Rust [`RustHeaderValue`] together with the GIL token so it can be
+/// converted into a typed Python [`HeaderValue`] (the `Raw` variant needs the
+/// token to allocate its `PyBytes` payload).
+struct RustHeaderValueRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value RustHeaderValue,
+}
+
+impl<'py> TryFrom<RustHeaderValueRef<'py, '_>> for Bound<'py, HeaderValue> {
+ type Error = PyErr;
+
+ fn try_from(value: RustHeaderValueRef<'py, '_>) -> PyResult<Self> {
+ let RustHeaderValueRef { py, value } = value;
+ let value = match value.kind() {
+ HeaderKind::Raw => HeaderValue::Raw {
+ value: PyBytes::new(py,
value.as_raw().map_err(to_value_error)?).unbind(),
+ },
+ HeaderKind::String => HeaderValue::String {
+ value: value.as_str().map_err(to_value_error)?.to_string(),
+ },
+ HeaderKind::Bool => HeaderValue::Bool {
+ value: value.as_bool().map_err(to_value_error)?,
+ },
+ HeaderKind::Int8 => HeaderValue::Int8 {
+ value: value.as_int8().map_err(to_value_error)?,
+ },
+ HeaderKind::Int16 => HeaderValue::Int16 {
+ value: value.as_int16().map_err(to_value_error)?,
+ },
+ HeaderKind::Int32 => HeaderValue::Int32 {
+ value: value.as_int32().map_err(to_value_error)?,
+ },
+ HeaderKind::Int64 => HeaderValue::Int64 {
+ value: value.as_int64().map_err(to_value_error)?,
+ },
+ HeaderKind::Int128 => HeaderValue::Int128 {
+ value: value.as_int128().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint8 => HeaderValue::UnsignedInt8 {
+ value: value.as_uint8().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint16 => HeaderValue::UnsignedInt16 {
+ value: value.as_uint16().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint32 => HeaderValue::UnsignedInt32 {
+ value: value.as_uint32().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint64 => HeaderValue::UnsignedInt64 {
+ value: value.as_uint64().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint128 => HeaderValue::UnsignedInt128 {
+ value: value.as_uint128().map_err(to_value_error)?,
+ },
+ HeaderKind::Float32 => HeaderValue::Float32 {
+ value: value.as_float32().map_err(to_value_error)?,
+ },
+ HeaderKind::Float64 => HeaderValue::Float64 {
+ value: value.as_float64().map_err(to_value_error)?,
+ },
+ };
+ value.into_pyobject(py)
+ }
+}
+
+impl HeaderValue {
+ fn identity(&self, py: Python<'_>) -> PyResult<HeaderIdentity> {
+ match self {
+ HeaderValue::Raw { value } =>
Ok(HeaderIdentity::raw(value.extract::<Vec<u8>>(py)?)),
+ HeaderValue::String { value } =>
Ok(HeaderIdentity::string(value.as_bytes().to_vec())),
+ HeaderValue::Bool { value } => Ok(HeaderIdentity::bool(*value)),
+ HeaderValue::Int8 { value } => Ok(HeaderIdentity::int8(*value)),
+ HeaderValue::Int16 { value } => Ok(HeaderIdentity::int16(*value)),
+ HeaderValue::Int32 { value } => Ok(HeaderIdentity::int32(*value)),
+ HeaderValue::Int64 { value } => Ok(HeaderIdentity::int64(*value)),
+ HeaderValue::Int128 { value } =>
Ok(HeaderIdentity::int128(*value)),
+ HeaderValue::UnsignedInt8 { value } =>
Ok(HeaderIdentity::uint8(*value)),
+ HeaderValue::UnsignedInt16 { value } =>
Ok(HeaderIdentity::uint16(*value)),
+ HeaderValue::UnsignedInt32 { value } =>
Ok(HeaderIdentity::uint32(*value)),
+ HeaderValue::UnsignedInt64 { value } =>
Ok(HeaderIdentity::uint64(*value)),
+ HeaderValue::UnsignedInt128 { value } =>
Ok(HeaderIdentity::uint128(*value)),
+ HeaderValue::Float32 { value } =>
Ok(HeaderIdentity::float32(*value)),
+ HeaderValue::Float64 { value } =>
Ok(HeaderIdentity::float64(*value)),
+ }
+ }
+
+ fn repr(&self, py: Python<'_>) -> PyResult<String> {
+ match self {
+ HeaderValue::Raw { value } => Ok(format!(
+ "HeaderValue.Raw({})",
+ value.bind(py).repr()?.extract::<String>()?
+ )),
+ HeaderValue::String { value } =>
Ok(format!("HeaderValue.String({value:?})")),
+ HeaderValue::Bool { value } =>
Ok(format!("HeaderValue.Bool({value})")),
+ HeaderValue::Int8 { value } =>
Ok(format!("HeaderValue.Int8({value})")),
+ HeaderValue::Int16 { value } =>
Ok(format!("HeaderValue.Int16({value})")),
+ HeaderValue::Int32 { value } =>
Ok(format!("HeaderValue.Int32({value})")),
+ HeaderValue::Int64 { value } =>
Ok(format!("HeaderValue.Int64({value})")),
+ HeaderValue::Int128 { value } =>
Ok(format!("HeaderValue.Int128({value})")),
+ HeaderValue::UnsignedInt8 { value } =>
Ok(format!("HeaderValue.UnsignedInt8({value})")),
+ HeaderValue::UnsignedInt16 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt16({value})"))
+ }
+ HeaderValue::UnsignedInt32 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt32({value})"))
+ }
+ HeaderValue::UnsignedInt64 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt64({value})"))
+ }
+ HeaderValue::UnsignedInt128 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt128({value})"))
+ }
+ HeaderValue::Float32 { value } =>
Ok(format!("HeaderValue.Float32({value:?})")),
+ HeaderValue::Float64 { value } =>
Ok(format!("HeaderValue.Float64({value:?})")),
+ }
+ }
+}
+
+#[derive(PartialEq, Eq, Hash)]
+struct HeaderIdentity {
+ kind: u8,
+ value: Vec<u8>,
+}
+
+impl HeaderIdentity {
+ fn raw(value: Vec<u8>) -> Self {
+ Self {
+ kind: HeaderKind::Raw.as_code(),
+ value,
+ }
+ }
+
+ fn string(value: Vec<u8>) -> Self {
+ Self {
+ kind: HeaderKind::String.as_code(),
+ value,
+ }
+ }
+
+ fn bool(value: bool) -> Self {
+ Self {
+ kind: HeaderKind::Bool.as_code(),
+ value: vec![u8::from(value)],
+ }
+ }
+
+ fn int8(value: i8) -> Self {
+ Self {
+ kind: HeaderKind::Int8.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int16(value: i16) -> Self {
+ Self {
+ kind: HeaderKind::Int16.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int32(value: i32) -> Self {
+ Self {
+ kind: HeaderKind::Int32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int64(value: i64) -> Self {
+ Self {
+ kind: HeaderKind::Int64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int128(value: i128) -> Self {
+ Self {
+ kind: HeaderKind::Int128.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint8(value: u8) -> Self {
+ Self {
+ kind: HeaderKind::Uint8.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint16(value: u16) -> Self {
+ Self {
+ kind: HeaderKind::Uint16.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint32(value: u32) -> Self {
+ Self {
+ kind: HeaderKind::Uint32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint64(value: u64) -> Self {
+ Self {
+ kind: HeaderKind::Uint64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint128(value: u128) -> Self {
+ Self {
+ kind: HeaderKind::Uint128.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn float32(value: f32) -> Self {
+ Self {
+ kind: HeaderKind::Float32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn float64(value: f64) -> Self {
+ Self {
+ kind: HeaderKind::Float64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+}
+
+pub(crate) fn py_user_headers_to_rust(
+ py: Python<'_>,
+ headers: &Bound<'_, PyDict>,
+) -> PyResult<RustUserHeaders> {
+ // Each key/value pair is converted independently: typed `HeaderKey` /
+ // `HeaderValue` are used as-is, while plain Python scalars are converted
on
+ // a best-effort basis.
+ let mut rust_headers = BTreeMap::new();
+ for (key, value) in headers.iter() {
+ let key = py_header_key_to_rust(py, &key)?;
+ let value = py_header_value_to_rust(py, &value)?;
+ if rust_headers.insert(key, value).is_some() {
+ return Err(PyValueError::new_err(
+ "Duplicate user header key: each header key must be unique",
+ ));
+ }
+ }
+ Ok(rust_headers)
+}
+
+pub(crate) fn rust_user_headers_to_py<'a>(
+ py: Python<'a>,
+ headers: RustUserHeaders,
+) -> PyResult<Bound<'a, UserHeaders>> {
+ // Always expose the explicitly typed dict[HeaderKey, HeaderValue] so that
+ // no wire-type information is silently dropped. Callers who prefer the
+ // convenient plain form opt in through `UserHeaders.to_scalar_dict`.
+ let result = Bound::new(py, UserHeaders)?;
+ let mapping = result.as_any();
+ for (key, value) in headers {
+ let key = Bound::<HeaderKey>::try_from(RustHeaderKeyRef { py, value:
&key })?;
+ let value = Bound::<HeaderValue>::try_from(RustHeaderValueRef { py,
value: &value })?;
+ mapping.set_item(key, value)?;
+ }
+ Ok(result)
+}
+
+/// User headers dictionary returned by `ReceiveMessage.user_headers`.
+///
+/// This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping
+/// operations work) that additionally exposes `to_scalar_dict` for the
convenient
+/// scalar form.
+#[gen_stub_pyclass]
+#[pyclass(extends=PyDict)]
+pub struct UserHeaders;
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl UserHeaders {
+ /// Wraps a mapping so its entries gain the `to_scalar_dict` helper.
+ ///
+ /// Accepts a dict whose keys and values can each independently be
+ /// `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool |
+ /// int | float`). The inherited `dict` initializer copies the provided
+ /// mapping.
+ #[new]
+ #[pyo3(signature = (mapping=None))]
+ pub fn new(
+ #[gen_stub(override_type(type_repr = "dict | None"))] mapping:
Option<&Bound<'_, PyAny>>,
+ ) -> Self {
+ let _ = mapping;
+ UserHeaders
+ }
+
+ /// Converts these headers into the convenient plain dictionary form.
+ ///
+ /// Returns an error if two distinct typed keys map to the same plain
+ /// Python scalar (e.g., `UnsignedInt8(1)` and `UnsignedInt16(1)` both
+ /// become `int(1)`), or if a stored field cannot be decoded.
+ #[gen_stub(override_return_type(
+ type_repr = "dict[str | bytes | bool | int | float, str | bytes | bool
| int | float]"
+ ))]
+ pub fn to_scalar_dict<'a>(slf: &Bound<'a, Self>) -> PyResult<Bound<'a,
PyDict>> {
+ let py = slf.py();
+ let dict = slf.as_any().cast::<PyDict>()?;
+ let result = PyDict::new(py);
+ for (key, value) in dict.iter() {
+ let plain_key = py_any_to_plain::<HeaderKey>(py, &key)?;
+ let plain_value = py_any_to_plain::<HeaderValue>(py, &value)?;
+ if result.contains(&plain_key)? {
+ return Err(PyValueError::new_err(
+ "Distinct typed header keys produce the same plain Python
scalar; this conversion is lossy and cannot proceed",
+ ));
+ }
+ result.set_item(plain_key, plain_value)?;
+ }
+ Ok(result)
+ }
+}
+
+fn py_any_to_plain<'py, T: ToPlainScalar + PyClass>(
+ py: Python<'py>,
+ any: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyAny>> {
+ if let Ok(header) = any.extract::<PyRef<'_, T>>() {
+ return Bound::<PyAny>::try_from(HeaderToPlainRef {
+ py,
+ value: &*header,
+ });
+ }
+ Ok(any.clone())
Review Comment:
Right now, since `UserHeaders` is a
publicly constructible and mutable `dict`, this is possible:
```python
UserHeaders({object(): object()}).to_scalar_dict()
```
The resulting dictionary can contain arbitrary objects despite the documented
and generated scalar-only return type.
Validate entries during construction and mutation, or reject unsupported
objects inside `to_scalar_dict()`. Add tests for unsupported keys and values
inserted through both the constructor and normal dictionary mutation.
##########
foreign/python/src/user_headers.rs:
##########
@@ -0,0 +1,897 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::BTreeMap;
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+
+use iggy::prelude::{
+ HeaderField, HeaderKey as RustHeaderKey, HeaderKind, HeaderValue as
RustHeaderValue,
+};
+use pyo3::PyClass;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::pyclass::CompareOp;
+use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyString};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum,
gen_stub_pymethods};
+
+type RustUserHeaders = BTreeMap<RustHeaderKey, RustHeaderValue>;
+
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+/// Typed key for an Iggy user header.
+///
+/// Use these constructors when the header key must preserve an explicit
+/// wire type instead of using the common string-key dictionary form.
+pub enum HeaderKey {
+ /// Raw bytes key. The byte length must be 1..=255.
+ Raw { value: Py<PyBytes> },
+ /// UTF-8 string key. The encoded byte length must be 1..=255.
+ String { value: String },
+ /// Boolean key.
+ Bool { value: bool },
+ /// Signed 8-bit integer key.
+ Int8 { value: i8 },
+ /// Signed 16-bit integer key.
+ Int16 { value: i16 },
+ /// Signed 32-bit integer key.
+ Int32 { value: i32 },
+ /// Signed 64-bit integer key.
+ Int64 { value: i64 },
+ /// Signed 128-bit integer key.
+ Int128 { value: i128 },
+ /// Unsigned 8-bit integer key.
+ UnsignedInt8 { value: u8 },
+ /// Unsigned 16-bit integer key.
+ UnsignedInt16 { value: u16 },
+ /// Unsigned 32-bit integer key.
+ UnsignedInt32 { value: u32 },
+ /// Unsigned 64-bit integer key.
+ UnsignedInt64 { value: u64 },
+ /// Unsigned 128-bit integer key.
+ UnsignedInt128 { value: u128 },
+ /// 32-bit floating point key.
+ Float32 { value: f32 },
+ /// 64-bit floating point key.
+ Float64 { value: f64 },
+}
+
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+/// Typed value for an Iggy user header.
+///
+/// Use these constructors when the header value must preserve an explicit
+/// wire type instead of using the common Python scalar dictionary form.
+pub enum HeaderValue {
+ /// Raw bytes value. The byte length must be 1..=255.
+ Raw { value: Py<PyBytes> },
+ /// UTF-8 string value. The encoded byte length must be 1..=255.
+ String { value: String },
+ /// Boolean value.
+ Bool { value: bool },
+ /// Signed 8-bit integer value.
+ Int8 { value: i8 },
+ /// Signed 16-bit integer value.
+ Int16 { value: i16 },
+ /// Signed 32-bit integer value.
+ Int32 { value: i32 },
+ /// Signed 64-bit integer value.
+ Int64 { value: i64 },
+ /// Signed 128-bit integer value.
+ Int128 { value: i128 },
+ /// Unsigned 8-bit integer value.
+ UnsignedInt8 { value: u8 },
+ /// Unsigned 16-bit integer value.
+ UnsignedInt16 { value: u16 },
+ /// Unsigned 32-bit integer value.
+ UnsignedInt32 { value: u32 },
+ /// Unsigned 64-bit integer value.
+ UnsignedInt64 { value: u64 },
+ /// Unsigned 128-bit integer value.
+ UnsignedInt128 { value: u128 },
+ /// 32-bit floating point value.
+ Float32 { value: f32 },
+ /// 64-bit floating point value.
+ Float64 { value: f64 },
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl HeaderKey {
+ pub fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
+ Ok(py_hash(header_identity_hash(self.identity(py)?)))
+ }
+
+ pub fn __richcmp__(
+ &self,
+ py: Python<'_>,
+ other: &Bound<'_, PyAny>,
+ op: CompareOp,
+ ) -> PyResult<Py<PyAny>> {
+ header_key_richcmp(py, self.identity(py)?, other, op)
+ }
+
+ pub fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+ self.repr(py)
+ }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl HeaderValue {
+ pub fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
+ Ok(py_hash(header_identity_hash(self.identity(py)?)))
+ }
+
+ pub fn __richcmp__(
+ &self,
+ py: Python<'_>,
+ other: &Bound<'_, PyAny>,
+ op: CompareOp,
+ ) -> PyResult<Py<PyAny>> {
+ header_value_richcmp(py, self.identity(py)?, other, op)
+ }
+
+ pub fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
+ self.repr(py)
+ }
+}
+
+/// Borrows a typed Python [`HeaderKey`] together with the GIL token so it can
+/// be converted into its Rust representation (the `Raw` variant needs the
token
+/// to read its `PyBytes` payload).
+struct PyHeaderKeyRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value HeaderKey,
+}
+
+impl TryFrom<PyHeaderKeyRef<'_, '_>> for RustHeaderKey {
+ type Error = PyErr;
+
+ fn try_from(key: PyHeaderKeyRef<'_, '_>) -> PyResult<Self> {
+ let PyHeaderKeyRef { py, value } = key;
+ match value {
+ HeaderKey::Raw { value } => {
+
RustHeaderKey::try_from(value.extract::<Vec<u8>>(py)?).map_err(to_value_error)
+ }
+ HeaderKey::String { value } => {
+ RustHeaderKey::try_from(value.as_str()).map_err(to_value_error)
+ }
+ HeaderKey::Bool { value } => Ok((*value).into()),
+ HeaderKey::Int8 { value } => Ok((*value).into()),
+ HeaderKey::Int16 { value } => Ok((*value).into()),
+ HeaderKey::Int32 { value } => Ok((*value).into()),
+ HeaderKey::Int64 { value } => Ok((*value).into()),
+ HeaderKey::Int128 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt8 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt16 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt32 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt64 { value } => Ok((*value).into()),
+ HeaderKey::UnsignedInt128 { value } => Ok((*value).into()),
+ HeaderKey::Float32 { value } => checked_float32(*value),
+ HeaderKey::Float64 { value } => Ok((*value).into()),
+ }
+ }
+}
+
+/// Borrows a Rust [`RustHeaderKey`] together with the GIL token so it can be
+/// converted into a typed Python [`HeaderKey`] (the `Raw` variant needs the
+/// token to allocate its `PyBytes` payload).
+struct RustHeaderKeyRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value RustHeaderKey,
+}
+
+impl<'py> TryFrom<RustHeaderKeyRef<'py, '_>> for Bound<'py, HeaderKey> {
+ type Error = PyErr;
+
+ fn try_from(key: RustHeaderKeyRef<'py, '_>) -> PyResult<Self> {
+ let RustHeaderKeyRef { py, value } = key;
+ let key = match value.kind() {
+ HeaderKind::Raw => HeaderKey::Raw {
+ value: PyBytes::new(py,
value.as_raw().map_err(to_value_error)?).unbind(),
+ },
+ HeaderKind::String => HeaderKey::String {
+ value: value.as_str().map_err(to_value_error)?.to_string(),
+ },
+ HeaderKind::Bool => HeaderKey::Bool {
+ value: value.as_bool().map_err(to_value_error)?,
+ },
+ HeaderKind::Int8 => HeaderKey::Int8 {
+ value: value.as_int8().map_err(to_value_error)?,
+ },
+ HeaderKind::Int16 => HeaderKey::Int16 {
+ value: value.as_int16().map_err(to_value_error)?,
+ },
+ HeaderKind::Int32 => HeaderKey::Int32 {
+ value: value.as_int32().map_err(to_value_error)?,
+ },
+ HeaderKind::Int64 => HeaderKey::Int64 {
+ value: value.as_int64().map_err(to_value_error)?,
+ },
+ HeaderKind::Int128 => HeaderKey::Int128 {
+ value: value.as_int128().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint8 => HeaderKey::UnsignedInt8 {
+ value: value.as_uint8().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint16 => HeaderKey::UnsignedInt16 {
+ value: value.as_uint16().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint32 => HeaderKey::UnsignedInt32 {
+ value: value.as_uint32().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint64 => HeaderKey::UnsignedInt64 {
+ value: value.as_uint64().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint128 => HeaderKey::UnsignedInt128 {
+ value: value.as_uint128().map_err(to_value_error)?,
+ },
+ HeaderKind::Float32 => HeaderKey::Float32 {
+ value: value.as_float32().map_err(to_value_error)?,
+ },
+ HeaderKind::Float64 => HeaderKey::Float64 {
+ value: value.as_float64().map_err(to_value_error)?,
+ },
+ };
+ key.into_pyobject(py)
+ }
+}
+
+trait ToPlainScalar {
+ fn to_plain_scalar<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py,
PyAny>>;
+}
+
+macro_rules! impl_to_plain_scalar {
+ ($ty:ty) => {
+ impl ToPlainScalar for $ty {
+ fn to_plain_scalar<'py>(&self, py: Python<'py>) ->
PyResult<Bound<'py, PyAny>> {
+ match self {
+ Self::Raw { value } =>
Ok(value.bind(py).clone().into_any()),
+ Self::String { value } =>
Ok(value.clone().into_pyobject(py)?.into_any()),
+ Self::Bool { value } =>
Ok((*value).into_pyobject(py)?.to_owned().into_any()),
+ Self::Int8 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int16 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int32 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Int128 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt8 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt16 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt32 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::UnsignedInt128 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ Self::Float32 { value } =>
Ok(f64::from(*value).into_pyobject(py)?.into_any()),
+ Self::Float64 { value } =>
Ok((*value).into_pyobject(py)?.into_any()),
+ }
+ }
+ }
+ };
+}
+
+impl_to_plain_scalar!(HeaderKey);
+impl_to_plain_scalar!(HeaderValue);
+
+/// Borrows a typed header ([`HeaderKey`] or [`HeaderValue`]) together with the
+/// GIL token for direct conversion to a plain Python scalar.
+struct HeaderToPlainRef<'py, 'value, T: ToPlainScalar> {
+ py: Python<'py>,
+ value: &'value T,
+}
+
+impl<'py, T: ToPlainScalar> TryFrom<HeaderToPlainRef<'py, '_, T>> for
Bound<'py, PyAny> {
+ type Error = PyErr;
+
+ fn try_from(inner: HeaderToPlainRef<'py, '_, T>) -> PyResult<Self> {
+ let HeaderToPlainRef { py, value } = inner;
+ value.to_plain_scalar(py)
+ }
+}
+
+impl HeaderKey {
+ fn identity(&self, py: Python<'_>) -> PyResult<HeaderIdentity> {
+ match self {
+ HeaderKey::Raw { value } =>
Ok(HeaderIdentity::raw(value.extract::<Vec<u8>>(py)?)),
+ HeaderKey::String { value } =>
Ok(HeaderIdentity::string(value.as_bytes().to_vec())),
+ HeaderKey::Bool { value } => Ok(HeaderIdentity::bool(*value)),
+ HeaderKey::Int8 { value } => Ok(HeaderIdentity::int8(*value)),
+ HeaderKey::Int16 { value } => Ok(HeaderIdentity::int16(*value)),
+ HeaderKey::Int32 { value } => Ok(HeaderIdentity::int32(*value)),
+ HeaderKey::Int64 { value } => Ok(HeaderIdentity::int64(*value)),
+ HeaderKey::Int128 { value } => Ok(HeaderIdentity::int128(*value)),
+ HeaderKey::UnsignedInt8 { value } =>
Ok(HeaderIdentity::uint8(*value)),
+ HeaderKey::UnsignedInt16 { value } =>
Ok(HeaderIdentity::uint16(*value)),
+ HeaderKey::UnsignedInt32 { value } =>
Ok(HeaderIdentity::uint32(*value)),
+ HeaderKey::UnsignedInt64 { value } =>
Ok(HeaderIdentity::uint64(*value)),
+ HeaderKey::UnsignedInt128 { value } =>
Ok(HeaderIdentity::uint128(*value)),
+ HeaderKey::Float32 { value } =>
Ok(HeaderIdentity::float32(*value)),
+ HeaderKey::Float64 { value } =>
Ok(HeaderIdentity::float64(*value)),
+ }
+ }
+
+ fn repr(&self, py: Python<'_>) -> PyResult<String> {
+ match self {
+ HeaderKey::Raw { value } => Ok(format!(
+ "HeaderKey.Raw({})",
+ value.bind(py).repr()?.extract::<String>()?
+ )),
+ HeaderKey::String { value } =>
Ok(format!("HeaderKey.String({value:?})")),
+ HeaderKey::Bool { value } =>
Ok(format!("HeaderKey.Bool({value})")),
+ HeaderKey::Int8 { value } =>
Ok(format!("HeaderKey.Int8({value})")),
+ HeaderKey::Int16 { value } =>
Ok(format!("HeaderKey.Int16({value})")),
+ HeaderKey::Int32 { value } =>
Ok(format!("HeaderKey.Int32({value})")),
+ HeaderKey::Int64 { value } =>
Ok(format!("HeaderKey.Int64({value})")),
+ HeaderKey::Int128 { value } =>
Ok(format!("HeaderKey.Int128({value})")),
+ HeaderKey::UnsignedInt8 { value } =>
Ok(format!("HeaderKey.UnsignedInt8({value})")),
+ HeaderKey::UnsignedInt16 { value } =>
Ok(format!("HeaderKey.UnsignedInt16({value})")),
+ HeaderKey::UnsignedInt32 { value } =>
Ok(format!("HeaderKey.UnsignedInt32({value})")),
+ HeaderKey::UnsignedInt64 { value } =>
Ok(format!("HeaderKey.UnsignedInt64({value})")),
+ HeaderKey::UnsignedInt128 { value } =>
Ok(format!("HeaderKey.UnsignedInt128({value})")),
+ HeaderKey::Float32 { value } =>
Ok(format!("HeaderKey.Float32({value:?})")),
+ HeaderKey::Float64 { value } =>
Ok(format!("HeaderKey.Float64({value:?})")),
+ }
+ }
+}
+
+/// Borrows a typed Python [`HeaderValue`] together with the GIL token so it
can
+/// be converted into its Rust representation (the `Raw` variant needs the
token
+/// to read its `PyBytes` payload).
+struct PyHeaderValueRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value HeaderValue,
+}
+
+impl TryFrom<PyHeaderValueRef<'_, '_>> for RustHeaderValue {
+ type Error = PyErr;
+
+ fn try_from(value: PyHeaderValueRef<'_, '_>) -> PyResult<Self> {
+ let PyHeaderValueRef { py, value } = value;
+ match value {
+ HeaderValue::Raw { value } => {
+
RustHeaderValue::try_from(value.extract::<Vec<u8>>(py)?).map_err(to_value_error)
+ }
+ HeaderValue::String { value } => {
+
RustHeaderValue::try_from(value.as_str()).map_err(to_value_error)
+ }
+ HeaderValue::Bool { value } => Ok((*value).into()),
+ HeaderValue::Int8 { value } => Ok((*value).into()),
+ HeaderValue::Int16 { value } => Ok((*value).into()),
+ HeaderValue::Int32 { value } => Ok((*value).into()),
+ HeaderValue::Int64 { value } => Ok((*value).into()),
+ HeaderValue::Int128 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt8 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt16 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt32 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt64 { value } => Ok((*value).into()),
+ HeaderValue::UnsignedInt128 { value } => Ok((*value).into()),
+ HeaderValue::Float32 { value } => checked_float32(*value),
+ HeaderValue::Float64 { value } => Ok((*value).into()),
+ }
+ }
+}
+
+/// Borrows a Rust [`RustHeaderValue`] together with the GIL token so it can be
+/// converted into a typed Python [`HeaderValue`] (the `Raw` variant needs the
+/// token to allocate its `PyBytes` payload).
+struct RustHeaderValueRef<'py, 'value> {
+ py: Python<'py>,
+ value: &'value RustHeaderValue,
+}
+
+impl<'py> TryFrom<RustHeaderValueRef<'py, '_>> for Bound<'py, HeaderValue> {
+ type Error = PyErr;
+
+ fn try_from(value: RustHeaderValueRef<'py, '_>) -> PyResult<Self> {
+ let RustHeaderValueRef { py, value } = value;
+ let value = match value.kind() {
+ HeaderKind::Raw => HeaderValue::Raw {
+ value: PyBytes::new(py,
value.as_raw().map_err(to_value_error)?).unbind(),
+ },
+ HeaderKind::String => HeaderValue::String {
+ value: value.as_str().map_err(to_value_error)?.to_string(),
+ },
+ HeaderKind::Bool => HeaderValue::Bool {
+ value: value.as_bool().map_err(to_value_error)?,
+ },
+ HeaderKind::Int8 => HeaderValue::Int8 {
+ value: value.as_int8().map_err(to_value_error)?,
+ },
+ HeaderKind::Int16 => HeaderValue::Int16 {
+ value: value.as_int16().map_err(to_value_error)?,
+ },
+ HeaderKind::Int32 => HeaderValue::Int32 {
+ value: value.as_int32().map_err(to_value_error)?,
+ },
+ HeaderKind::Int64 => HeaderValue::Int64 {
+ value: value.as_int64().map_err(to_value_error)?,
+ },
+ HeaderKind::Int128 => HeaderValue::Int128 {
+ value: value.as_int128().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint8 => HeaderValue::UnsignedInt8 {
+ value: value.as_uint8().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint16 => HeaderValue::UnsignedInt16 {
+ value: value.as_uint16().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint32 => HeaderValue::UnsignedInt32 {
+ value: value.as_uint32().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint64 => HeaderValue::UnsignedInt64 {
+ value: value.as_uint64().map_err(to_value_error)?,
+ },
+ HeaderKind::Uint128 => HeaderValue::UnsignedInt128 {
+ value: value.as_uint128().map_err(to_value_error)?,
+ },
+ HeaderKind::Float32 => HeaderValue::Float32 {
+ value: value.as_float32().map_err(to_value_error)?,
+ },
+ HeaderKind::Float64 => HeaderValue::Float64 {
+ value: value.as_float64().map_err(to_value_error)?,
+ },
+ };
+ value.into_pyobject(py)
+ }
+}
+
+impl HeaderValue {
+ fn identity(&self, py: Python<'_>) -> PyResult<HeaderIdentity> {
+ match self {
+ HeaderValue::Raw { value } =>
Ok(HeaderIdentity::raw(value.extract::<Vec<u8>>(py)?)),
+ HeaderValue::String { value } =>
Ok(HeaderIdentity::string(value.as_bytes().to_vec())),
+ HeaderValue::Bool { value } => Ok(HeaderIdentity::bool(*value)),
+ HeaderValue::Int8 { value } => Ok(HeaderIdentity::int8(*value)),
+ HeaderValue::Int16 { value } => Ok(HeaderIdentity::int16(*value)),
+ HeaderValue::Int32 { value } => Ok(HeaderIdentity::int32(*value)),
+ HeaderValue::Int64 { value } => Ok(HeaderIdentity::int64(*value)),
+ HeaderValue::Int128 { value } =>
Ok(HeaderIdentity::int128(*value)),
+ HeaderValue::UnsignedInt8 { value } =>
Ok(HeaderIdentity::uint8(*value)),
+ HeaderValue::UnsignedInt16 { value } =>
Ok(HeaderIdentity::uint16(*value)),
+ HeaderValue::UnsignedInt32 { value } =>
Ok(HeaderIdentity::uint32(*value)),
+ HeaderValue::UnsignedInt64 { value } =>
Ok(HeaderIdentity::uint64(*value)),
+ HeaderValue::UnsignedInt128 { value } =>
Ok(HeaderIdentity::uint128(*value)),
+ HeaderValue::Float32 { value } =>
Ok(HeaderIdentity::float32(*value)),
+ HeaderValue::Float64 { value } =>
Ok(HeaderIdentity::float64(*value)),
+ }
+ }
+
+ fn repr(&self, py: Python<'_>) -> PyResult<String> {
+ match self {
+ HeaderValue::Raw { value } => Ok(format!(
+ "HeaderValue.Raw({})",
+ value.bind(py).repr()?.extract::<String>()?
+ )),
+ HeaderValue::String { value } =>
Ok(format!("HeaderValue.String({value:?})")),
+ HeaderValue::Bool { value } =>
Ok(format!("HeaderValue.Bool({value})")),
+ HeaderValue::Int8 { value } =>
Ok(format!("HeaderValue.Int8({value})")),
+ HeaderValue::Int16 { value } =>
Ok(format!("HeaderValue.Int16({value})")),
+ HeaderValue::Int32 { value } =>
Ok(format!("HeaderValue.Int32({value})")),
+ HeaderValue::Int64 { value } =>
Ok(format!("HeaderValue.Int64({value})")),
+ HeaderValue::Int128 { value } =>
Ok(format!("HeaderValue.Int128({value})")),
+ HeaderValue::UnsignedInt8 { value } =>
Ok(format!("HeaderValue.UnsignedInt8({value})")),
+ HeaderValue::UnsignedInt16 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt16({value})"))
+ }
+ HeaderValue::UnsignedInt32 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt32({value})"))
+ }
+ HeaderValue::UnsignedInt64 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt64({value})"))
+ }
+ HeaderValue::UnsignedInt128 { value } => {
+ Ok(format!("HeaderValue.UnsignedInt128({value})"))
+ }
+ HeaderValue::Float32 { value } =>
Ok(format!("HeaderValue.Float32({value:?})")),
+ HeaderValue::Float64 { value } =>
Ok(format!("HeaderValue.Float64({value:?})")),
+ }
+ }
+}
+
+#[derive(PartialEq, Eq, Hash)]
+struct HeaderIdentity {
+ kind: u8,
+ value: Vec<u8>,
+}
+
+impl HeaderIdentity {
+ fn raw(value: Vec<u8>) -> Self {
+ Self {
+ kind: HeaderKind::Raw.as_code(),
+ value,
+ }
+ }
+
+ fn string(value: Vec<u8>) -> Self {
+ Self {
+ kind: HeaderKind::String.as_code(),
+ value,
+ }
+ }
+
+ fn bool(value: bool) -> Self {
+ Self {
+ kind: HeaderKind::Bool.as_code(),
+ value: vec![u8::from(value)],
+ }
+ }
+
+ fn int8(value: i8) -> Self {
+ Self {
+ kind: HeaderKind::Int8.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int16(value: i16) -> Self {
+ Self {
+ kind: HeaderKind::Int16.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int32(value: i32) -> Self {
+ Self {
+ kind: HeaderKind::Int32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int64(value: i64) -> Self {
+ Self {
+ kind: HeaderKind::Int64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn int128(value: i128) -> Self {
+ Self {
+ kind: HeaderKind::Int128.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint8(value: u8) -> Self {
+ Self {
+ kind: HeaderKind::Uint8.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint16(value: u16) -> Self {
+ Self {
+ kind: HeaderKind::Uint16.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint32(value: u32) -> Self {
+ Self {
+ kind: HeaderKind::Uint32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint64(value: u64) -> Self {
+ Self {
+ kind: HeaderKind::Uint64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn uint128(value: u128) -> Self {
+ Self {
+ kind: HeaderKind::Uint128.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn float32(value: f32) -> Self {
+ Self {
+ kind: HeaderKind::Float32.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+
+ fn float64(value: f64) -> Self {
+ Self {
+ kind: HeaderKind::Float64.as_code(),
+ value: value.to_le_bytes().to_vec(),
+ }
+ }
+}
+
+pub(crate) fn py_user_headers_to_rust(
+ py: Python<'_>,
+ headers: &Bound<'_, PyDict>,
+) -> PyResult<RustUserHeaders> {
+ // Each key/value pair is converted independently: typed `HeaderKey` /
+ // `HeaderValue` are used as-is, while plain Python scalars are converted
on
+ // a best-effort basis.
+ let mut rust_headers = BTreeMap::new();
+ for (key, value) in headers.iter() {
+ let key = py_header_key_to_rust(py, &key)?;
+ let value = py_header_value_to_rust(py, &value)?;
+ if rust_headers.insert(key, value).is_some() {
+ return Err(PyValueError::new_err(
+ "Duplicate user header key: each header key must be unique",
+ ));
+ }
+ }
+ Ok(rust_headers)
+}
+
+pub(crate) fn rust_user_headers_to_py<'a>(
+ py: Python<'a>,
+ headers: RustUserHeaders,
+) -> PyResult<Bound<'a, UserHeaders>> {
+ // Always expose the explicitly typed dict[HeaderKey, HeaderValue] so that
+ // no wire-type information is silently dropped. Callers who prefer the
+ // convenient plain form opt in through `UserHeaders.to_scalar_dict`.
+ let result = Bound::new(py, UserHeaders)?;
+ let mapping = result.as_any();
+ for (key, value) in headers {
+ let key = Bound::<HeaderKey>::try_from(RustHeaderKeyRef { py, value:
&key })?;
+ let value = Bound::<HeaderValue>::try_from(RustHeaderValueRef { py,
value: &value })?;
+ mapping.set_item(key, value)?;
+ }
+ Ok(result)
+}
+
+/// User headers dictionary returned by `ReceiveMessage.user_headers`.
+///
+/// This is a regular `dict[HeaderKey, HeaderValue]` (so all mapping
+/// operations work) that additionally exposes `to_scalar_dict` for the
convenient
+/// scalar form.
+#[gen_stub_pyclass]
+#[pyclass(extends=PyDict)]
+pub struct UserHeaders;
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl UserHeaders {
+ /// Wraps a mapping so its entries gain the `to_scalar_dict` helper.
+ ///
+ /// Accepts a dict whose keys and values can each independently be
+ /// `HeaderKey`/`HeaderValue` or a plain scalar (`str | bytes | bool |
+ /// int | float`). The inherited `dict` initializer copies the provided
+ /// mapping.
+ #[new]
+ #[pyo3(signature = (mapping=None))]
+ pub fn new(
+ #[gen_stub(override_type(type_repr = "dict | None"))] mapping:
Option<&Bound<'_, PyAny>>,
+ ) -> Self {
+ let _ = mapping;
+ UserHeaders
+ }
+
+ /// Converts these headers into the convenient plain dictionary form.
+ ///
+ /// Returns an error if two distinct typed keys map to the same plain
+ /// Python scalar (e.g., `UnsignedInt8(1)` and `UnsignedInt16(1)` both
+ /// become `int(1)`), or if a stored field cannot be decoded.
+ #[gen_stub(override_return_type(
+ type_repr = "dict[str | bytes | bool | int | float, str | bytes | bool
| int | float]"
+ ))]
+ pub fn to_scalar_dict<'a>(slf: &Bound<'a, Self>) -> PyResult<Bound<'a,
PyDict>> {
+ let py = slf.py();
+ let dict = slf.as_any().cast::<PyDict>()?;
+ let result = PyDict::new(py);
+ for (key, value) in dict.iter() {
+ let plain_key = py_any_to_plain::<HeaderKey>(py, &key)?;
+ let plain_value = py_any_to_plain::<HeaderValue>(py, &value)?;
+ if result.contains(&plain_key)? {
+ return Err(PyValueError::new_err(
+ "Distinct typed header keys produce the same plain Python
scalar; this conversion is lossy and cannot proceed",
+ ));
+ }
+ result.set_item(plain_key, plain_value)?;
+ }
+ Ok(result)
+ }
+}
+
+fn py_any_to_plain<'py, T: ToPlainScalar + PyClass>(
+ py: Python<'py>,
+ any: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyAny>> {
+ if let Ok(header) = any.extract::<PyRef<'_, T>>() {
+ return Bound::<PyAny>::try_from(HeaderToPlainRef {
+ py,
+ value: &*header,
+ });
+ }
+ Ok(any.clone())
+}
+
+fn py_header_key_to_rust(py: Python<'_>, key: &Bound<'_, PyAny>) ->
PyResult<RustHeaderKey> {
+ // A typed `HeaderKey` already carries an explicit wire type.
+ if let Ok(header_key) = key.extract::<PyRef<'_, HeaderKey>>() {
+ return RustHeaderKey::try_from(PyHeaderKeyRef {
+ py,
+ value: &header_key,
+ });
+ }
+
+ // Otherwise best-effort convert the plain Python scalar.
+ plain_to_header_field(key)?.ok_or_else(|| {
+ PyValueError::new_err("User header keys must be str, bytes, bool, int,
float, or HeaderKey")
+ })
+}
+
+fn py_header_value_to_rust(py: Python<'_>, value: &Bound<'_, PyAny>) ->
PyResult<RustHeaderValue> {
+ // A typed `HeaderValue` already carries an explicit wire type.
+ if let Ok(header_value) = value.extract::<PyRef<'_, HeaderValue>>() {
+ return RustHeaderValue::try_from(PyHeaderValueRef {
+ py,
+ value: &header_value,
+ });
+ }
+
+ // Otherwise best-effort convert the plain Python scalar.
+ plain_to_header_field(value)?.ok_or_else(|| {
+ PyValueError::new_err(
+ "User header values must be str, bytes, bool, int, float, or
HeaderValue",
+ )
+ })
+}
+
+/// Converts a plain Python scalar into a header field (used for both keys and
+/// values). Returns `Ok(None)` when the object is not a supported plain type
so
+/// the caller can raise a message naming keys vs values; genuine range/length
+/// errors are returned as `Err`.
+fn plain_to_header_field<T>(obj: &Bound<'_, PyAny>) ->
PyResult<Option<HeaderField<T>>> {
+ // `bool` must be checked before `int` since Python bools are also ints.
+ if obj.is_instance_of::<PyBool>() {
+ return Ok(Some(obj.extract::<bool>()?.into()));
+ }
+ if obj.is_instance_of::<PyInt>() {
+ return int_to_header_field(obj).map(Some);
+ }
+ if obj.is_instance_of::<PyFloat>() {
+ return Ok(Some(float_to_header_field(obj.extract::<f64>()?)));
+ }
+ if obj.is_instance_of::<PyString>() {
+ return HeaderField::try_from(obj.extract::<String>()?)
+ .map(Some)
+ .map_err(to_value_error);
+ }
+ if obj.is_instance_of::<PyBytes>() {
+ return HeaderField::try_from(obj.extract::<Vec<u8>>()?)
+ .map(Some)
+ .map_err(to_value_error);
+ }
+ Ok(None)
+}
+
+/// Picks the smallest header integer kind that represents the value exactly:
+/// non-negative values use the narrowest unsigned kind, negatives the
narrowest
+/// signed kind. Values beyond the 128-bit range are rejected.
+fn int_to_header_field<T>(obj: &Bound<'_, PyAny>) -> PyResult<HeaderField<T>> {
+ if let Ok(value) = obj.extract::<u8>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<u16>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<u32>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<u64>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<u128>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<i8>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<i16>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<i32>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<i64>() {
+ return Ok(value.into());
+ }
+ if let Ok(value) = obj.extract::<i128>() {
+ return Ok(value.into());
+ }
+ Err(PyValueError::new_err(
+ "User header int values must fit within the 128-bit range",
+ ))
+}
+
+/// Uses `Float32` when the value is exactly representable as an f32, otherwise
+/// `Float64`, so the plain float is always stored losslessly.
+fn float_to_header_field<T>(value: f64) -> HeaderField<T> {
+ let narrowed = value as f32;
+ if f64::from(narrowed) == value {
+ narrowed.into()
+ } else {
+ value.into()
+ }
+}
Review Comment:
This accepts plain `float("inf")` and stores it as `Float32`, because
narrowing infinity to `f32` still compares equal after widening. In contrast,
`checked_float32` reject the equivalent typed
`HeaderValue.Float32(float("inf"))`.
This also means a `Float32` infinity or NaN received from a Rust client can
be
represented by Python but cannot be sent again. `Float64` and plain NaN
follow
different paths again.
Let's reject non-finite headers everywhere in plain, Float32, and Float64
paths.
Add tests for positive and negative infinity and NaN across plain, typed, and
receive-then-resend paths.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
