graphrecords_python/querying/
argument.rs

1use crate::querying::{PyOperand, value_conversion::ValueConversion};
2use graphrecords_query::{
3    Mask, Scalar,
4    dynamic::{DynArgumentSource, DynInvokeArgument, DynValue},
5    registry::ValueDescriptor,
6};
7use pyo3::{
8    exceptions::PyTypeError,
9    prelude::*,
10    types::{PyBytes, PyString},
11};
12
13#[pyclass(frozen)]
14pub struct PyArgument(DynArgumentSource);
15
16impl PyArgument {
17    pub(super) const fn new(source: DynArgumentSource) -> Self {
18        Self(source)
19    }
20
21    pub(super) const fn source(&self) -> &DynArgumentSource {
22        &self.0
23    }
24}
25
26impl PyOperand {
27    pub(super) fn mask_argument(object: &Bound<'_, PyAny>) -> PyResult<DynInvokeArgument> {
28        if let Some(source) = Self::operand_source(object) {
29            return Ok(DynInvokeArgument::Source(source));
30        }
31
32        Ok(DynInvokeArgument::Source(DynArgumentSource::mask(
33            object.extract()?,
34        )))
35    }
36
37    pub(super) fn scalar_argument(object: &Bound<'_, PyAny>) -> PyResult<DynInvokeArgument> {
38        if let Some(source) = Self::operand_source(object) {
39            return Ok(DynInvokeArgument::Source(source));
40        }
41        let value = DynValue::from_python(object, &ValueDescriptor::value::<Scalar>())?;
42
43        Ok(DynInvokeArgument::Source(DynArgumentSource::value(value)))
44    }
45
46    pub(super) fn operand_argument(object: &Bound<'_, PyAny>) -> PyResult<DynInvokeArgument> {
47        object
48            .cast::<Self>()
49            .map(|operand| DynInvokeArgument::Operand(operand.get().operand().clone()))
50            .map_err(|_| PyTypeError::new_err("expected an operand argument"))
51    }
52
53    pub(super) fn value_argument(&self, object: &Bound<'_, PyAny>) -> PyResult<DynInvokeArgument> {
54        if let Some(source) = Self::operand_source(object) {
55            return Ok(DynInvokeArgument::Source(source));
56        }
57        let descriptor = self.operand().descriptor().lane_shape().value();
58
59        if descriptor.domain().is::<Mask>() {
60            return Ok(DynInvokeArgument::Source(DynArgumentSource::mask(
61                object.extract()?,
62            )));
63        }
64
65        DynValue::from_python(object, descriptor)
66            .map(DynArgumentSource::value)
67            .map(DynInvokeArgument::Source)
68    }
69
70    pub(super) fn set_argument(&self, values: &Bound<'_, PyAny>) -> PyResult<DynInvokeArgument> {
71        if let Some(source) = Self::operand_source(values) {
72            return Ok(DynInvokeArgument::Source(source));
73        }
74        if values.is_instance_of::<PyString>() || values.is_instance_of::<PyBytes>() {
75            return Err(PyTypeError::new_err(
76                "expected a sequence of values; `str` and `bytes` are single values",
77            ));
78        }
79        let descriptor = self.operand().descriptor().lane_shape().value();
80
81        if descriptor.domain().is::<Mask>() {
82            let values = values
83                .try_iter()?
84                .map(|value| value?.extract())
85                .collect::<PyResult<Vec<_>>>()?;
86
87            return Ok(DynInvokeArgument::Source(DynArgumentSource::mask_values(
88                values,
89            )));
90        }
91        let values = values
92            .try_iter()?
93            .map(|value| DynValue::from_python(&value?, descriptor))
94            .collect::<PyResult<Vec<_>>>()?;
95
96        Ok(DynInvokeArgument::Source(DynArgumentSource::values(values)))
97    }
98
99    pub(super) fn dropping_argument(&self) -> PyArgument {
100        PyArgument::new(DynArgumentSource::drop_missing(self.operand().clone()))
101    }
102
103    pub(super) fn replacing_argument(
104        &self,
105        replacement: &Bound<'_, PyAny>,
106    ) -> PyResult<PyArgument> {
107        if replacement.cast::<PyArgument>().is_ok() {
108            return Err(PyTypeError::new_err(
109                "an `on_missing` argument cannot itself be used as a replacement",
110            ));
111        }
112        let source = self.operand().clone();
113
114        if let Ok(operand) = replacement.cast::<Self>() {
115            let replacement = operand.get().operand().clone();
116
117            return Ok(PyArgument::new(
118                DynArgumentSource::replace_missing_with_operand(source, replacement),
119            ));
120        }
121        let descriptor = self.operand().descriptor().lane_shape().value();
122
123        if descriptor.domain().is::<Mask>() {
124            return Ok(PyArgument::new(
125                DynArgumentSource::replace_missing_with_mask(source, replacement.extract()?),
126            ));
127        }
128
129        DynValue::from_python(replacement, descriptor).map(|value| {
130            PyArgument::new(DynArgumentSource::replace_missing_with_value(source, value))
131        })
132    }
133
134    fn operand_source(object: &Bound<'_, PyAny>) -> Option<DynArgumentSource> {
135        if let Ok(argument) = object.cast::<PyArgument>() {
136            return Some(argument.get().source().clone());
137        }
138
139        object
140            .cast::<Self>()
141            .ok()
142            .map(|operand| DynArgumentSource::operand(operand.get().operand().clone()))
143    }
144}
graphrecords_python/querying/argument.rs

graphrecords_python/querying/
argument.rs
