py_mcai_worker_sdk 1.1.3

AMQP Worker to listen and provide trait to process message
#[cfg(feature = "media")]
use crate::python::GenericStreamDescriptor;
use mcai_worker_sdk::prelude::*;
use pyo3::{prelude::*, types::*};
use std::{env, fs};

pub(crate) fn read_python_file() -> String {
  let filename = env::var("PYTHON_WORKER_FILENAME").unwrap_or_else(|_| "worker.py".to_string());

  fs::read_to_string(&filename)
    .unwrap_or_else(|_| panic!("unable to open and read file: {}", filename))
}

pub(crate) fn get_python_module(gil: &GILGuard) -> Result<(Python, &PyModule)> {
  let python_file_content = read_python_file();
  let py = gil.python();
  let python_module = PyModule::from_code(py, &python_file_content, "worker.py", "worker")
    .map_err(|error| {
      MessageError::RuntimeError(format!(
        "unable to create the python module: {}",
        py_err_to_string(py, error)
      ))
    })?;
  Ok((py, python_module))
}

pub(crate) fn get_string_from_module(method: &str) -> String {
  let gil = Python::acquire_gil();
  let (_py, python_module) = get_python_module(&gil).unwrap();

  let response: String = python_module
    .getattr(method)
    .unwrap_or_else(|_| panic!("unable to get {} attribute in your module", method))
    .call0()
    .unwrap_or_else(|_| panic!("unable to call {} in your module", method))
    .extract()
    .unwrap_or_else(|_| panic!("unable to found a return value for {} function", method));

  response
}

pub(crate) fn call_module_function<'a>(
  py: Python,
  python_module: &'a PyModule,
  function_name: &'a str,
  args: impl IntoPy<Py<PyTuple>>,
) -> std::result::Result<&'a PyAny, String> {
  python_module
    .getattr(function_name)
    .unwrap_or_else(|_| panic!("unable to get {} attribute in your module", function_name))
    .call1(args)
    .map_err(move |error| {
      let ptraceback = &error.ptraceback(py);
      let stacktrace = ptraceback
        .as_ref()
        .map(|tb| {
          let traceback = py.import("traceback").unwrap();
          let locals = [("traceback", traceback)].into_py_dict(py);

          locals.set_item("tb", tb).unwrap();

          py.eval("traceback.format_tb(tb)", None, Some(locals))
            .expect("Unknown python error, unable to get the stacktrace")
            .to_string()
        })
        .unwrap_or_else(|| "Unknown python error, no stackstrace".to_string());

      let error_msg = py_err_to_string(py, error);

      format!("{}\n\nStacktrace:\n{}", error_msg, stacktrace)
    })
}

pub(crate) fn get_parameters() -> Vec<WorkerParameter> {
  let gil = Python::acquire_gil();
  let (py, python_module) = get_python_module(&gil).unwrap();

  let method = "get_parameters";
  let response = python_module
    .getattr(method)
    .unwrap_or_else(|_| panic!("unable to get {} attribute in your module", method))
    .call0()
    .unwrap_or_else(|_| panic!("unable to call {} in your module", method))
    .downcast::<PyList>()
    .unwrap();

  let mut parameters = vec![];

  for item in response.iter() {
    let object = item.downcast::<PyDict>().expect("not a python dict");

    let label = object
      .get_item("label")
      .expect("missing label in parameter")
      .to_string();
    let identifier = object
      .get_item("identifier")
      .expect("missing identifier in parameter")
      .to_string();

    let kind_list = object
      .get_item("kind")
      .expect("missing kind in parameter")
      .downcast::<PyList>()
      .unwrap();

    let mut parameter_types = vec![];

    for kind in kind_list.iter() {
      let value = kind
        .downcast::<PyString>()
        .expect("not a python string")
        .to_string();
      let parameter_type: WorkerParameterType =
        serde_json::from_str(&format!("{:?}", value)).unwrap();
      parameter_types.push(parameter_type);
    }

    let required = object
      .get_item("required")
      .unwrap_or_else(|| PyBool::new(py, false).as_ref())
      .is_true()
      .unwrap();

    parameters.push(WorkerParameter {
      label,
      identifier,
      kind: parameter_types,
      required,
    });
  }

  parameters
}

pub fn py_err_to_string(py: Python, error: PyErr) -> String {
  let locals = [("error", error)].into_py_dict(py);

  py.eval("repr(error)", None, Some(locals))
    .expect("Unknown python error, unable to get the error message")
    .to_string()
}

pub fn get_destination_paths(response: &PyAny) -> Option<Vec<String>> {
  if response.is_none() {
    return None;
  }

  response
    .downcast::<PyDict>()
    .map(|object| {
      object
        .get_item("destination_paths")
        .map(|response_paths| {
          response_paths
            .downcast::<PyList>()
            .map(|path_list| {
              let destination_paths = path_list
                .iter()
                .map(|item| item.downcast::<PyString>())
                .filter(|downcast| downcast.is_ok())
                .map(|value| value.unwrap().to_string())
                .collect();

              Some(destination_paths)
            })
            .unwrap_or(None)
        })
        .flatten()
    })
    .unwrap_or(None)
}

#[cfg(feature = "media")]
pub fn get_stream_descriptors(response: &PyAny) -> Result<Vec<StreamDescriptor>> {
  response
    .downcast::<PyList>()
    .map(|py_list| {
      py_list
        .iter()
        .map(|value| value.extract::<GenericStreamDescriptor>())
        .filter(|extracted| extracted.is_ok())
        .map(|extracted| get_stream_descriptor(extracted.unwrap()))
        .collect()
    })
    .map_err(|e| {
      MessageError::RuntimeError(format!(
        "unable to access init_process(..) python response: {:?}",
        e
      ))
    })
}

#[cfg(feature = "media")]
fn get_stream_descriptor(generic_stream_descriptor: GenericStreamDescriptor) -> StreamDescriptor {
  if generic_stream_descriptor.stream_type.is_video() {
    let filters = generic_stream_descriptor
      .filters
      .iter()
      .cloned()
      .map(VideoFilter::Generic)
      .collect();
    StreamDescriptor::new_video(generic_stream_descriptor.index as usize, filters)
  } else if generic_stream_descriptor.stream_type.is_audio() {
    let filters = generic_stream_descriptor
      .filters
      .iter()
      .cloned()
      .map(AudioFilter::Generic)
      .collect();
    StreamDescriptor::new_audio(generic_stream_descriptor.index as usize, filters)
  } else {
    StreamDescriptor::new_data(generic_stream_descriptor.index as usize)
  }
}