diff --git a/src/ast.rs b/src/ast.rs index 5050a84..a11d353 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -210,6 +210,14 @@ impl Sub for EvaluateResult { } impl Expr<'_> { + pub fn get_representation<'a>(self) -> Result<&'a str, Box> { + use Expr::*; + match self { + Dielectric(_) => Ok("eps"), + Index(_) => Ok("n"), + _ => Err("Not a valid expression".into()), + } + } pub fn evaluate<'a>( &self, params: &mut ExprParams<'a>, @@ -244,9 +252,7 @@ impl Expr<'_> { _ => Err(MissingParameter::new(key).into()), }, Dielectric(ref expr) => expr.evaluate(params), - Index(ref expr) => Ok(expr - .evaluate(params)? - .power(EvaluateResult::Number(Complex64::from(2.)))), + Index(ref expr) => expr.evaluate(params), Sum(ref expr) => { let mut params_vec = Vec::new(); for (key, val) in params.rep_params.iter() { diff --git a/src/formula_parser.lalrpop b/src/formula_parser.lalrpop index b1aa531..bd1a1b8 100644 --- a/src/formula_parser.lalrpop +++ b/src/formula_parser.lalrpop @@ -82,5 +82,5 @@ Num: f64 = { }; Var: &'input str = { - r"[a-zA-Z][a-zA-Z_0-0]*" + r"[a-zA-Z][a-zA-Z_0-9]*" }; diff --git a/src/lib.rs b/src/lib.rs index 9594651..8c62a7b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -102,16 +102,29 @@ fn parse<'a>( rep_params, sum_params: None, })? { - EvaluateResult::Number(num) => Ok(Array1::from_elem( - x_axis_values.len(), - Complex64::from(num), - )), - EvaluateResult::Array(arr) => Ok(arr) + EvaluateResult::Number(num) => { + Ok(Array1::from_elem(x_axis_values.len(), Complex64::from(num))) + } + EvaluateResult::Array(arr) => Ok(arr), } } #[pymodule] fn formula_dispersion(_py: Python, m: &PyModule) -> PyResult<()> { + #[pyfn(m)] + #[pyo3(name = "get_representation")] + fn get_representation_py<'py>(formula: &str) -> PyResult<&'py str> { + let ast = match parse_ast(formula) { + Ok(ast) => ast, + Err(err) => return Err(PyErr::new::(err.to_string())), + }; + + match ast.get_representation() { + Ok(repr) => Ok(repr), + Err(err) => Err(PyErr::new::(err.to_string())), + } + } + #[pyfn(m)] #[pyo3(name = "parse")] fn parse_py<'py>(