use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use crate::core::stop_condition::StopCondition;
#[pyclass(name = "StopCondition")]
#[derive(Clone)]
pub struct PyStopCondition {
pub(crate) inner: StopCondition,
}
impl PyStopCondition {
pub(crate) fn into_inner(self) -> StopCondition {
self.inner
}
}
#[pymethods]
impl PyStopCondition {
#[staticmethod]
fn max_rows(n: u64) -> PyResult<Self> {
if n == 0 {
return Err(PyValueError::new_err("n must be > 0"));
}
Ok(Self {
inner: StopCondition::MaxRows(n),
})
}
#[staticmethod]
fn max_bytes(n: u64) -> PyResult<Self> {
if n == 0 {
return Err(PyValueError::new_err("n must be > 0"));
}
Ok(Self {
inner: StopCondition::MaxBytes(n),
})
}
#[staticmethod]
fn schema_stable(consecutive_stable_rows: u64) -> PyResult<Self> {
if consecutive_stable_rows == 0 {
return Err(PyValueError::new_err("consecutive_stable_rows must be > 0"));
}
Ok(Self {
inner: StopCondition::SchemaStable {
consecutive_stable_rows,
},
})
}
#[staticmethod]
fn confidence_threshold(threshold: f64) -> PyResult<Self> {
if !(0.0..=1.0).contains(&threshold) {
return Err(PyValueError::new_err(
"threshold must be between 0.0 and 1.0",
));
}
Ok(Self {
inner: StopCondition::ConfidenceThreshold(threshold),
})
}
#[staticmethod]
fn memory_pressure(threshold: f64) -> PyResult<Self> {
if !(0.0..=1.0).contains(&threshold) {
return Err(PyValueError::new_err(
"threshold must be between 0.0 and 1.0",
));
}
Ok(Self {
inner: StopCondition::MemoryPressure(threshold),
})
}
#[staticmethod]
fn never() -> Self {
Self {
inner: StopCondition::Never,
}
}
#[staticmethod]
fn schema_inference() -> Self {
Self {
inner: StopCondition::schema_inference(),
}
}
#[staticmethod]
fn quality_sample() -> Self {
Self {
inner: StopCondition::quality_sample(),
}
}
fn __or__(&self, other: &PyStopCondition) -> Self {
let mut conditions = Vec::new();
match &self.inner {
StopCondition::Any(v) => conditions.extend(v.clone()),
c => conditions.push(c.clone()),
}
match &other.inner {
StopCondition::Any(v) => conditions.extend(v.clone()),
c => conditions.push(c.clone()),
}
Self {
inner: StopCondition::Any(conditions),
}
}
fn __and__(&self, other: &PyStopCondition) -> Self {
let mut conditions = Vec::new();
match &self.inner {
StopCondition::All(v) => conditions.extend(v.clone()),
c => conditions.push(c.clone()),
}
match &other.inner {
StopCondition::All(v) => conditions.extend(v.clone()),
c => conditions.push(c.clone()),
}
Self {
inner: StopCondition::All(conditions),
}
}
fn __repr__(&self) -> String {
format!("StopCondition({:?})", self.inner)
}
}