py_mcai_worker_sdk 1.1.3

#[macro_use]
extern crate serde_derive;

use crate::helpers::{call_module_function, get_python_module, get_string_from_module};
use crate::python::CallbackHandle;
use crate::{
  helpers::get_destination_paths,
  parameters::{build_parameters, PythonWorkerParameters},
};
#[cfg(feature = "media")]
use crate::{
  helpers::get_stream_descriptors, media::PyEbuTtmlLive, python::StreamDescriptorHandler,
};
use mcai_worker_sdk::prelude::*;
use pyo3::prelude::*;
#[cfg(feature = "media")]
use std::{
  convert::TryInto,
  ops::Deref,
  sync::{mpsc::Sender, Arc, Mutex},
};

mod helpers;
#[cfg(feature = "media")]
mod media;
mod parameters;
mod python;
#[cfg(test)]
mod tests;

#[derive(Clone, Debug, Default)]
struct PythonWorkerEvent {
  #[cfg(feature = "media")]
  result: Option<Arc<Mutex<Sender<ProcessResult>>>>,
}

impl McaiWorker<PythonWorkerParameters> for PythonWorkerEvent {
  fn get_name(&self) -> String {
    get_string_from_module("get_name")
  }

  fn get_short_description(&self) -> String {
    get_string_from_module("get_short_description")
  }

  fn get_description(&self) -> String {
    get_string_from_module("get_description")
  }

  fn get_version(&self) -> Version {
    Version::parse(&get_string_from_module("get_version"))
      .expect("unable to parse version (please use SemVer format)")
  }

  fn init(&mut self) -> Result<()> {
    let gil = Python::acquire_gil();
    let (py, python_module) = get_python_module(&gil)?;

    let optional_init_function_name = "init";

    if python_module.getattr(optional_init_function_name).is_ok() {
      let _result = call_module_function(py, python_module, optional_init_function_name, ())
        .map_err(MessageError::ParameterValueError)?;
    } else {
      info!(
        "No optional '{}' function to call.",
        optional_init_function_name
      );
    }

    Ok(())
  }

  #[cfg(feature = "media")]
  fn init_process(
    &mut self,
    parameters: PythonWorkerParameters,
    format_context: Arc<Mutex<FormatContext>>,
    result: Arc<Mutex<Sender<ProcessResult>>>,
  ) -> Result<Vec<StreamDescriptor>> {
    self.result = Some(result);

    let gil = Python::acquire_gil();
    let (py, python_module) = get_python_module(&gil)?;

    let context = media::FormatContext::from(format_context);
    let list_of_parameters = build_parameters(parameters, py)?;

    let handler = StreamDescriptorHandler {};

    let response = call_module_function(
      py,
      python_module,
      "init_process",
      (handler, context, list_of_parameters),
    )
    .map_err(MessageError::ParameterValueError)?;
    get_stream_descriptors(response)
  }

  #[cfg(feature = "media")]
  fn process_frame(
    &mut self,
    job_result: JobResult,
    stream_index: usize,
    process_frame: ProcessFrame,
  ) -> Result<ProcessResult> {
    let gil = Python::acquire_gil();
    let (py, python_module) = get_python_module(&gil)?;

    match &process_frame {
      ProcessFrame::AudioVideo(frame) => {
        let media_frame: media::Frame = frame.try_into()?;

        let response = call_module_function(
          py,
          python_module,
          "process_frame",
          (&job_result.get_str_job_id(), stream_index, media_frame),
        )
        .map_err(|error_message| {
          let result = job_result
            .with_status(JobStatus::Error)
            .with_message(&error_message);
          MessageError::ProcessingError(result)
        })?;

        Ok(ProcessResult::new_json(&response.to_string()))
      }
      ProcessFrame::EbuTtmlLive(ebu_ttml_live) => {
        let ttml_content: PyEbuTtmlLive = ebu_ttml_live.deref().clone().into();

        let response = call_module_function(
          py,
          python_module,
          "process_ebu_ttml_live",
          (&job_result.get_str_job_id(), stream_index, ttml_content),
        )
        .map_err(|error_message| {
          let result = job_result
            .with_status(JobStatus::Error)
            .with_message(&error_message);
          MessageError::ProcessingError(result)
        })?;

        Ok(ProcessResult::new_json(&response.to_string()))
      }
      ProcessFrame::Json(_json) => Err(MessageError::NotImplemented()),
      ProcessFrame::Data(_data) => Err(MessageError::NotImplemented()),
    }
  }

  #[cfg(feature = "media")]
  fn ending_process(&mut self) -> Result<()> {
    let gil = Python::acquire_gil();
    let (py, python_module) = get_python_module(&gil)?;

    let _result = call_module_function(py, python_module, "ending_process", ())
      .map_err(MessageError::ParameterValueError)?;

    if let Some(result) = &self.result {
      result
        .lock()
        .unwrap()
        .send(ProcessResult::end_of_process())
        .unwrap();
    }

    Ok(())
  }

  fn process(
    &self,
    channel: Option<McaiChannel>,
    parameters: PythonWorkerParameters,
    mut job_result: JobResult,
  ) -> Result<JobResult> {
    let gil = Python::acquire_gil();
    let (py, python_module) = get_python_module(&gil)?;

    let list_of_parameters = build_parameters(parameters, py)?;

    let job_status = Arc::new(Mutex::new(None));

    let callback_handle = CallbackHandle {
      channel,
      job_id: job_result.get_job_id(),
      job_status: job_status.clone(),
    };

    let response = call_module_function(
      py,
      python_module,
      "process",
      (callback_handle, list_of_parameters, job_result.get_job_id()),
    )
    .map_err(|error_message| {
      let result = job_result
        .clone()
        .with_status(JobStatus::Error)
        .with_message(&error_message);
      MessageError::ProcessingError(result)
    })?;

    if let Some(mut destination_paths) = get_destination_paths(response) {
      job_result = job_result.with_destination_paths(&mut destination_paths);
    }

    let final_status = job_status
      .lock()
      .unwrap()
      .as_ref()
      .unwrap_or(&JobStatus::Completed)
      .clone();
    Ok(job_result.with_status(final_status))
  }
}

fn main() {
  start_worker(PythonWorkerEvent::default());
}