python-betterproto/betterproto-extras/src/merging.rs

use crate::{error::Result, py_any_extras::PyAnyExtras};
use indoc::indoc;
use prost_reflect::{
    prost_types::{Duration, Timestamp},
    DynamicMessage, MapKey, ReflectMessage, Value,
};
use pyo3::{
    sync::GILOnceCell,
    types::{IntoPyDict, PyBytes, PyModule},
    Py, PyAny, PyObject, Python, ToPyObject,
};

pub fn merge_msg_into_pyobj(obj: &PyAny, mut msg: DynamicMessage) -> Result<()> {
    for field in msg.take_fields() {
        let field_name = field.0.name();
        let proto_meta = obj.get_proto_meta()?;
        obj.setattr(
            field_name,
            map_field_value(field_name, field.1, proto_meta)?,
        )?;
    }

    let mut buf = vec![];
    for field in msg.unknown_fields() {
        field.encode(&mut buf);
    }
    if !buf.is_empty() {
        let mut unknown_fields = obj.getattr("_unknown_fields")?.extract::<Vec<u8>>()?;
        unknown_fields.append(&mut buf);
        obj.setattr("_unknown_fields", PyBytes::new(obj.py(), &unknown_fields))?;
    }

    obj.setattr("_serialized_on_wire", true)?;

    Ok(())
}

fn map_field_value(field_name: &str, field_value: Value, proto_meta: &PyAny) -> Result<PyObject> {
    let py = proto_meta.py();
    match field_value {
        Value::Bool(x) => Ok(x.to_object(py)),
        Value::Bytes(x) => Ok(PyBytes::new(py, &x).to_object(py)),
        Value::F32(x) => Ok(x.to_object(py)),
        Value::F64(x) => Ok(x.to_object(py)),
        Value::I32(x) => Ok(x.to_object(py)),
        Value::I64(x) => Ok(x.to_object(py)),
        Value::String(x) => Ok(x.to_object(py)),
        Value::U32(x) => Ok(x.to_object(py)),
        Value::U64(x) => Ok(x.to_object(py)),
        Value::Message(msg) => match msg.descriptor().full_name() {
            "google.protobuf.BoolValue" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_bool())
                .to_object(py)),
            "google.protobuf.DoubleValue" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_f64())
                .to_object(py)),
            "google.protobuf.FloatValue" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_f32())
                .to_object(py)),
            "google.protobuf.Int64Value" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_i64())
                .to_object(py)),
            "google.protobuf.UInt64Value" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_u64())
                .to_object(py)),
            "google.protobuf.Int32Value" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_i32())
                .to_object(py)),
            "google.protobuf.UInt32Value" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_u32())
                .to_object(py)),
            "google.protobuf.StringValue" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_str().map(|s| s.to_string()))
                .to_object(py)),
            "google.protobuf.BytesValue" => Ok(msg
                .get_field_by_number(1)
                .and_then(|val| val.as_bytes().map(|b| PyBytes::new(py, b)))
                .to_object(py)),
            "google.protobuf.Timestamp" => {
                let msg = msg.transcode_to::<Timestamp>()?;
                Ok(create_py_datetime(&msg, py))
            }
            "google.protobuf.Duration" => {
                let msg = msg.transcode_to::<Duration>()?;
                Ok(create_py_timedelta(&msg, py))
            }
            _ => {
                let obj = proto_meta.create_instance(field_name)?;
                merge_msg_into_pyobj(obj, msg)?;
                Ok(obj.to_object(py))
            }
        },
        Value::List(ls) => Ok(ls
            .into_iter()
            .map(|x| map_field_value(field_name, x, proto_meta))
            .collect::<Result<Vec<PyObject>>>()?
            .to_object(py)),
        Value::EnumNumber(x) => {
            let cls = proto_meta.get_class(field_name)?;
            Ok(cls.call1((x,))?.to_object(py))
        }
        Value::Map(map) => {
            let res: Result<Vec<_>> = map
                .into_iter()
                .map(|(k, v)| {
                    let key = map_key(k, py);
                    let val = map_field_value(&format!("{field_name}.value"), v, proto_meta)?;
                    Ok((key, val))
                })
                .collect();
            Ok(res?.into_py_dict(py).to_object(py))
        }
    }
}

fn map_key(key: MapKey, py: Python) -> PyObject {
    match key {
        MapKey::Bool(x) => x.to_object(py),
        MapKey::I32(x) => x.to_object(py),
        MapKey::I64(x) => x.to_object(py),
        MapKey::U32(x) => x.to_object(py),
        MapKey::U64(x) => x.to_object(py),
        MapKey::String(x) => x.to_object(py),
    }
}

fn create_py_datetime(ts: &Timestamp, py: Python) -> PyObject {
    static CONSTRUCTOR_CACHE: GILOnceCell<Py<PyAny>> = GILOnceCell::new();
    let constructor = CONSTRUCTOR_CACHE.get_or_init(py, || {
        let constructor = PyModule::from_code(
            py,
            indoc! {"
                from datetime import datetime, timezone
                
                def constructor(ts):
                    return datetime.fromtimestamp(ts, tz=timezone.utc)
            "},
            "",
            "",
        )
        .expect("This is a valid Python module")
        .getattr("constructor")
        .expect("Attribute exists");
        Py::from(constructor)
    });
    let ts = (ts.seconds as f64) + (ts.nanos as f64) / 1e9;
    constructor
        .call1(py, (ts,))
        .expect("static function will not fail")
}

fn create_py_timedelta(duration: &Duration, py: Python) -> PyObject {
    static CONSTRUCTOR_CACHE: GILOnceCell<Py<PyAny>> = GILOnceCell::new();
    let constructor = CONSTRUCTOR_CACHE.get_or_init(py, || {
        let constructor = PyModule::from_code(
            py,
            indoc! {"
                from datetime import timedelta
                
                def constructor(s, ms):
                    return timedelta(seconds=s, microseconds=ms)
            "},
            "",
            "",
        )
        .expect("This is a valid Python module")
        .getattr("constructor")
        .expect("Attribute exists");
        Py::from(constructor)
    });
    constructor
        .call1(py, (duration.seconds as f64, (duration.nanos as f64) / 1e3))
        .expect("static function will not fail")
}