Adds representation function

src/ast.rs

@@ -210,6 +210,14 @@ impl Sub for EvaluateResult {
 }
 
 impl Expr<'_> {
+    pub fn get_representation<'a>(self) -> Result<&'a str, Box<dyn error::Error>> {
+        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() {

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]*"
 };

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::<PyTypeError, _>(err.to_string())),
+        };
+
+        match ast.get_representation() {
+            Ok(repr) => Ok(repr),
+            Err(err) => Err(PyErr::new::<PyTypeError, _>(err.to_string())),
+        }
+    }
+
     #[pyfn(m)]
     #[pyo3(name = "parse")]
     fn parse_py<'py>(