diff --git a/tests/test_formula_evaluate.py b/tests/test_formula_evaluate.py index b1c2e9c..e03d8a1 100644 --- a/tests/test_formula_evaluate.py +++ b/tests/test_formula_evaluate.py @@ -225,6 +225,29 @@ def evaluate_expr(source: str, columns: dict[str, Col]) -> object: columns={"ra": Col("00 02 08.4")}, error=True, ), + EvalCase( + name="unit_from_string", + expression='unit("km")', + columns={}, + result_val=1.0, + result_unit=u.km, + ), + EvalCase( + name="unit_multiplied", + expression='col("float_col") * unit("Mpc")', + columns=_COLUMNS, + result_val=1.5, + result_unit=u.Mpc, + ), + EvalCase( + name="unit_composite", + expression='col("float_col") * unit("km/s")', + columns=_COLUMNS, + result_val=1.5, + result_unit=u.km / u.s, + ), + EvalCase(name="error_unknown_unit", expression='unit("not_a_unit")', columns={}, error=True), + EvalCase(name="error_unit_non_string", expression="unit(1)", columns={}, error=True), ] diff --git a/tests/test_formula_parse.py b/tests/test_formula_parse.py index e4cab6c..dd9b717 100644 --- a/tests/test_formula_parse.py +++ b/tests/test_formula_parse.py @@ -7,6 +7,7 @@ ('col("weird name")', {"weird name"}), ('sin(col("pa")) + pi', {"pa"}), ('to_deg(col("RAJ2000"))', {"RAJ2000"}), + ('col("dist") * unit("Mpc")', {"dist"}), ('3 * 10 ** col("logd25") * col("e_logd25") * arcsec', {"logd25", "e_logd25"}), ('"M " + col("id")', {"id"}), ("1 + 2", set()), diff --git a/tests/test_formula_tokens.py b/tests/test_formula_tokens.py index be02027..68b108e 100644 --- a/tests/test_formula_tokens.py +++ b/tests/test_formula_tokens.py @@ -4,7 +4,7 @@ def test_expression_tokens_include_language_names() -> None: labels = {token["label"] for token in expression_tokens()} - assert labels >= {"col", "sin", "cos", "str", "to_deg", "pi", "deg", "arcsec", "mag"} + assert labels >= {"col", "sin", "cos", "str", "to_deg", "unit", "pi", "deg", "arcsec", "mag"} def test_designation_form_marks_expression_widget() -> None: diff --git a/tests/test_formula_validate.py b/tests/test_formula_validate.py index ccc365c..c1695a1 100644 --- a/tests/test_formula_validate.py +++ b/tests/test_formula_validate.py @@ -6,6 +6,7 @@ def test_validate_expression_accepts_valid() -> None: assert validate_expression('to_deg(col("RAJ2000"))') == [] assert validate_expression("sin(pi) + 1.5 * deg") == [] + assert validate_expression('col("dist") * unit("Mpc")') == [] assert validate_expression("") == [] diff --git a/uploader/app/lib/formula/namespace.py b/uploader/app/lib/formula/namespace.py index 4160d86..8e080de 100644 --- a/uploader/app/lib/formula/namespace.py +++ b/uploader/app/lib/formula/namespace.py @@ -101,13 +101,29 @@ def _to_deg(value: object) -> u.Quantity: raise TypeError(f"to_deg() expected angle or coordinate string, got {type(value).__name__}") +def _unit(name: object) -> u.Quantity: + if not isinstance(name, str): + raise TypeError(f"unit() expected a unit name string, got {type(name).__name__}") + return 1 * u.Unit(name) + + COL_FUNCTION = FunctionDef("col", "Rawdata column", placeholder='"${1:name}"') FUNCTIONS: tuple[FunctionDef, ...] = ( FunctionDef("sin", "Sine (argument must be an angle)", np.sin), FunctionDef("cos", "Cosine (argument must be an angle)", np.cos), FunctionDef("str", "Convert to text", _formula_str), - FunctionDef("to_deg", "Convert to degrees; parses coordinate strings or angle quantities", _to_deg), + FunctionDef( + "to_deg", + 'Convert to degrees; e.g. "00 02 08.4" (hourangle), "+16 35 13" (deg), "00h02m08.4s"', + _to_deg, + ), + FunctionDef( + "unit", + 'Astropy unit from a name string; e.g. "Mpc", "km/s", "Jy"', + _unit, + placeholder='"${1:name}"', + ), ) @@ -182,5 +198,4 @@ def expression_syntax_help() -> str: - `180 * deg` - `"G"` - fills the column with a text "G" - Copy another column: `col("ra")` -- Sexagesimal coordinates: `to_deg(col("RAJ2000"))` - Mathematical expression: `3 * 10 ** col("logd25") * arcsec`"""