use fidius_core::python_descriptor::PythonInterfaceDescriptor;
use fidius_core::PluginError;
use pyo3::prelude::*;
use pyo3::types::{PyAnyMethods, PyBytes, PyDict, PyTuple};
use crate::error::pyerr_to_plugin_error;
use crate::value_bridge::{pyobject_to_value, value_to_pyobject};
#[derive(Debug, thiserror::Error)]
pub enum PythonCallError {
#[error("invalid method index {index} (interface has {count} method(s))")]
InvalidMethodIndex { index: usize, count: usize },
#[error(
"wire-mode mismatch on method '{method}': declared wire_raw={declared}, dispatcher used wire_raw={attempted}"
)]
WireModeMismatch {
method: &'static str,
declared: bool,
attempted: bool,
},
#[error("failed to decode typed input: {0}")]
InputDecode(String),
#[error("failed to encode typed output: {0}")]
OutputEncode(String),
#[error("plugin raised: [{}] {}", .0.code, .0.message)]
Plugin(PluginError),
}
#[derive(Debug)]
pub struct PythonPluginHandle {
descriptor: &'static PythonInterfaceDescriptor,
_module: Py<PyAny>,
method_callables: Vec<Py<PyAny>>,
}
impl PythonPluginHandle {
pub(crate) fn new(
descriptor: &'static PythonInterfaceDescriptor,
module: Py<PyAny>,
method_callables: Vec<Py<PyAny>>,
) -> Self {
Self {
descriptor,
_module: module,
method_callables,
}
}
pub fn descriptor(&self) -> &'static PythonInterfaceDescriptor {
self.descriptor
}
pub fn method_count(&self) -> usize {
self.descriptor.methods.len()
}
pub fn call_typed(
&self,
method_index: usize,
input_bincode: &[u8],
) -> Result<Vec<u8>, PythonCallError> {
self.call_typed_json(method_index, input_bincode)
}
pub fn call_typed_json(
&self,
method_index: usize,
input_json: &[u8],
) -> Result<Vec<u8>, PythonCallError> {
let method = self.lookup_method(method_index, false)?;
let input_value: serde_json::Value = serde_json::from_slice(input_json)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result_value = Python::with_gil(|py| -> Result<serde_json::Value, PythonCallError> {
let callable = method.callable.bind(py);
let py_args = build_call_args(py, &input_value)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result = callable
.call(py_args, None::<&Bound<'_, PyDict>>)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?;
pyobject_to_value(&result).map_err(|e| PythonCallError::OutputEncode(e.to_string()))
})?;
serde_json::to_vec(&result_value).map_err(|e| PythonCallError::OutputEncode(e.to_string()))
}
pub fn call_client_streaming_json(
&self,
method_index: usize,
items: Box<dyn Iterator<Item = serde_json::Value> + Send>,
args_json: &[u8],
) -> Result<Vec<u8>, PythonCallError> {
let method = self.lookup_method(method_index, false)?;
let args: serde_json::Value = serde_json::from_slice(args_json)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result_value = Python::with_gil(|py| -> Result<serde_json::Value, PythonCallError> {
let callable = method.callable.bind(py);
let stream = Py::new(
py,
HostFedStream {
items: std::sync::Mutex::new(items),
},
)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?
.into_bound(py)
.into_any();
let mut py_args: Vec<Bound<'_, PyAny>> = vec![stream];
match &args {
serde_json::Value::Array(a) => {
for v in a {
py_args.push(
value_to_pyobject(py, v)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?,
);
}
}
serde_json::Value::Null => {}
other => py_args.push(
value_to_pyobject(py, other)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?,
),
}
let args_tuple = PyTuple::new(py, py_args)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result = callable
.call(args_tuple, None::<&Bound<'_, PyDict>>)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?;
pyobject_to_value(&result).map_err(|e| PythonCallError::OutputEncode(e.to_string()))
})?;
serde_json::to_vec(&result_value).map_err(|e| PythonCallError::OutputEncode(e.to_string()))
}
pub fn call_bidi_streaming_start(
&self,
method_index: usize,
items: Box<dyn Iterator<Item = serde_json::Value> + Send>,
args_json: &[u8],
) -> Result<crate::stream::PythonStream, PythonCallError> {
let method = self.lookup_method(method_index, false)?;
let args: serde_json::Value = serde_json::from_slice(args_json)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
Python::with_gil(|py| {
let callable = method.callable.bind(py);
let stream = Py::new(
py,
HostFedStream {
items: std::sync::Mutex::new(items),
},
)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?
.into_bound(py)
.into_any();
let mut py_args: Vec<Bound<'_, PyAny>> = vec![stream];
match &args {
serde_json::Value::Array(a) => {
for v in a {
py_args.push(
value_to_pyobject(py, v)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?,
);
}
}
serde_json::Value::Null => {}
other => py_args.push(
value_to_pyobject(py, other)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?,
),
}
let args_tuple = PyTuple::new(py, py_args)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result = callable
.call(args_tuple, None::<&Bound<'_, PyDict>>)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?;
let iter = result.try_iter().map_err(|e| {
PythonCallError::OutputEncode(format!(
"bidirectional method must return an iterable/generator, got: {e}"
))
})?;
Ok(crate::stream::PythonStream::new(iter.into_any().unbind()))
})
}
pub fn call_streaming_start(
&self,
method_index: usize,
input_json: &[u8],
) -> Result<crate::stream::PythonStream, PythonCallError> {
let method = self.lookup_method(method_index, false)?;
let input_value: serde_json::Value = serde_json::from_slice(input_json)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
Python::with_gil(|py| {
let callable = method.callable.bind(py);
let py_args = build_call_args(py, &input_value)
.map_err(|e| PythonCallError::InputDecode(e.to_string()))?;
let result = callable
.call(py_args, None::<&Bound<'_, PyDict>>)
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?;
let iter = result.try_iter().map_err(|e| {
PythonCallError::OutputEncode(format!(
"streaming method must return an iterable/generator, got: {e}"
))
})?;
Ok(crate::stream::PythonStream::new(iter.into_any().unbind()))
})
}
pub fn call_raw(&self, method_index: usize, input: &[u8]) -> Result<Vec<u8>, PythonCallError> {
let method = self.lookup_method(method_index, true)?;
Python::with_gil(|py| {
let callable = method.callable.bind(py);
let arg = PyBytes::new(py, input);
let result = callable
.call1((arg,))
.map_err(|e| PythonCallError::Plugin(pyerr_to_plugin_error(e)))?;
let bytes: Vec<u8> = result.extract().map_err(|e| {
PythonCallError::OutputEncode(format!(
"raw method must return bytes/bytearray, got: {e}"
))
})?;
Ok(bytes)
})
}
fn lookup_method(
&self,
index: usize,
attempting_raw: bool,
) -> Result<MethodLookup<'_>, PythonCallError> {
if index >= self.method_callables.len() {
return Err(PythonCallError::InvalidMethodIndex {
index,
count: self.method_callables.len(),
});
}
let desc = &self.descriptor.methods[index];
if desc.wire_raw != attempting_raw {
return Err(PythonCallError::WireModeMismatch {
method: desc.name,
declared: desc.wire_raw,
attempted: attempting_raw,
});
}
Ok(MethodLookup {
callable: &self.method_callables[index],
})
}
}
struct MethodLookup<'a> {
callable: &'a Py<PyAny>,
}
fn build_call_args<'py>(
py: Python<'py>,
input: &serde_json::Value,
) -> PyResult<Bound<'py, PyTuple>> {
match input {
serde_json::Value::Array(items) => {
let py_items: Vec<Bound<'_, PyAny>> = items
.iter()
.map(|v| value_to_pyobject(py, v))
.collect::<PyResult<_>>()?;
PyTuple::new(py, py_items)
}
serde_json::Value::Null => PyTuple::new(py, Vec::<Bound<'_, PyAny>>::new()),
other => {
let pyobj = value_to_pyobject(py, other)?;
PyTuple::new(py, vec![pyobj])
}
}
}
#[pyclass]
struct HostFedStream {
items: std::sync::Mutex<Box<dyn Iterator<Item = serde_json::Value> + Send>>,
}
#[pymethods]
impl HostFedStream {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}
fn __next__(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
let next = self.items.lock().unwrap().next();
match next {
Some(v) => Ok(Some(value_to_pyobject(py, &v)?.unbind())),
None => Ok(None),
}
}
}