use numpy::{Element, PyArray1};
use pyo3::exceptions::PyTypeError;
use pyo3::prelude::pyclass;
use pyo3::types::PyDict;
use pyo3::{FromPyObject, PyAny, PyObject, PyResult};
use std::convert::TryInto;
use std::ffi::CString;
use std::os::raw::c_char;
pub enum Vectorized<'lt, T: Copy + Clone> {
Slice(&'lt [T]),
Scalar(T),
}
impl<'lt, T: Copy + Clone> Vectorized<'lt, T> {
#[allow(clippy::missing_safety_doc)]
#[inline]
pub unsafe fn get_unchecked(&self, index: usize) -> T {
match self {
Self::Scalar(res) => *res,
Self::Slice(values) => *values.get_unchecked(index),
}
}
}
impl<'lt, T: Copy + Clone + Element + FromPyObject<'lt> + 'lt> Vectorized<'lt, T> {
#[inline]
pub fn from_python(from: &'lt PyAny, name: &str, iterations: &mut usize) -> PyResult<Self> {
let res = if let Ok(val) = from.extract() {
Vectorized::Scalar(val)
} else if let Ok(values) = from.extract::<&PyArray1<T>>() {
match unsafe{values.as_slice()}.unwrap() {
[val] => Vectorized::Scalar(*val),
values if values.len() == *iterations => Vectorized::Slice(values),
values if *iterations == 1 => {
*iterations = values.len();
Vectorized::Slice(values)
},
values => return Err(pyo3::exceptions::PyTypeError::new_err(
format!(
"Arguments must have the same length or be scalars but '{}' has length {} while previous arguments had length {}",
name,
values.len(),
*iterations
)
)),
}
} else {
return Err(pyo3::exceptions::PyTypeError::new_err(format!(
"eval: Expected scalar {} value or 1d numpy {}-array for argument{}",
std::any::type_name::<T>(),
std::any::type_name::<T>(),
name,
)));
};
Ok(res)
}
#[inline]
pub fn from_dict(
dict: Option<&'lt PyDict>,
name: &str,
loc: &str,
iterations: &mut usize,
) -> PyResult<Self> {
if let Some(dict) = dict {
if let Some(val) = dict.get_item(name) {
return Vectorized::from_python(val, name, iterations);
}
}
Err(pyo3::exceptions::PyTypeError::new_err(format!(
"eval: Required argument '{}' is missing from {}",
name, loc
)))
}
}
impl<'lt> Vectorized<'lt, f64> {
#[inline]
pub fn from_dict_opt(
dict: Option<&'lt PyDict>,
name: &str,
iterations: &mut usize,
) -> PyResult<Self> {
if let Some(dict) = dict {
if let Some(val) = dict.get_item(name) {
return Vectorized::from_python(val, name, iterations);
}
}
Ok(Self::Scalar(0.0))
}
}
#[pyclass]
pub struct NonNumericParameter {
#[pyo3(get)]
pub name: &'static str,
#[pyo3(get)]
pub description: &'static str,
#[pyo3(get)]
pub unit: &'static str,
#[pyo3(get)]
pub group: &'static str,
#[pyo3(get)]
pub default: PyObject,
}
#[pyclass]
pub struct RealParameter {
#[pyo3(get)]
pub name: &'static str,
#[pyo3(get)]
pub description: &'static str,
#[pyo3(get)]
pub unit: &'static str,
#[pyo3(get)]
pub group: &'static str,
#[pyo3(get)]
pub default: f64,
#[pyo3(get)]
pub min: f64,
#[pyo3(get)]
pub max: f64,
#[pyo3(get)]
pub min_inclusive: bool,
#[pyo3(get)]
pub max_inclusive: bool,
}
#[pyclass]
pub struct IntegerParameter {
#[pyo3(get)]
pub name: &'static str,
#[pyo3(get)]
pub description: &'static str,
#[pyo3(get)]
pub unit: &'static str,
#[pyo3(get)]
pub group: &'static str,
#[pyo3(get)]
pub default: i64,
#[pyo3(get)]
pub min: i64,
#[pyo3(get)]
pub max: i64,
#[pyo3(get)]
pub min_inclusive: bool,
#[pyo3(get)]
pub max_inclusive: bool,
}
#[repr(transparent)]
pub struct PyFfiString(*mut c_char);
impl<'source> FromPyObject<'source> for PyFfiString {
fn extract(ob: &'source PyAny) -> PyResult<Self> {
let string: &str = ob.extract()?;
let cstr = CString::new(string).map_err(PyTypeError::new_err)?;
Ok(Self(cstr.into_raw()))
}
}
impl Drop for PyFfiString {
fn drop(&mut self) {
unsafe {
CString::from_raw(self.0);
}
}
}
pub trait FFI {
type FfiTy;
fn into_ffi(self) -> Self::FfiTy;
}
impl FFI for PyFfiString {
type FfiTy = *const c_char;
#[inline(always)]
fn into_ffi(self) -> Self::FfiTy {
self.0
}
}
impl FFI for bool {
type FfiTy = bool;
#[inline(always)]
fn into_ffi(self) -> Self::FfiTy {
self
}
}
impl FFI for f64 {
type FfiTy = f64;
#[inline(always)]
fn into_ffi(self) -> Self::FfiTy {
self
}
}
impl FFI for i64 {
type FfiTy = i64;
#[inline(always)]
fn into_ffi(self) -> Self::FfiTy {
self
}
}
#[repr(transparent)]
struct PyFfiPtr<A>(*mut A);
impl<A> Drop for PyFfiPtr<A> {
fn drop(&mut self) {
unsafe {
Box::from_raw(self.0);
}
}
}
#[pyclass]
pub struct OpVar {
#[pyo3(get)]
pub name: &'static str,
#[pyo3(get)]
pub description: &'static str,
#[pyo3(get)]
pub unit: &'static str,
}
macro_rules! array_ffi_impl(
($($len:expr),*) => {
$(
impl<'a, T: FromPyObject<'a> > FromPyObject<'a> for PyFfiPtr<[T; $len]> {
fn extract(ob: &'a PyAny) -> PyResult<Self> {
let data: Vec<T> = ob.extract()?;
let len = data.len();
let data: Box<[T; $len]> = data.into_boxed_slice().try_into().map_err(|_|PyTypeError::new_err(format!("Expected array of size {} but found {}",$len,len)))?;
Ok(Self(Box::into_raw(data)))
}
}
impl<T: FFI> FFI for PyFfiPtr<[T; $len]> {
type FfiTy = * const T::FfiTy;
fn into_ffi(self) -> Self::FfiTy {
self.0 as *const T::FfiTy
}
}
impl<T: FFI> FFI for [T; $len] {
type FfiTy = Self;
fn into_ffi(self) -> Self::FfiTy {
self
}
}
)*
}
);
array_ffi_impl!(
32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9,
8, 7, 6, 5, 4, 3, 2, 1, 0
);