add more functions to expression engine

.agents/skills/code-review/SKILL.md

@@ -119,3 +119,7 @@ An HTTP client fetches a URL taken from user/external input with no allowlist or
 ### command-injection
 
 subprocess/os.system/os.popen runs a shell with concatenated or formatted user/external input (shell=True, or a single string command). Only flag when the injected value is not a constant/whitelist.
+
+### pointless-wrapper
+
+A new or changed function or method only forwards to another callable with the same arguments and return value, adding no conversion, validation, defaulting, error handling, or other logic. Only flag when call sites could invoke the inner callable directly, the wrapper does not implement an interface, protocol, or abstract method, and it is not a public re-export of a private or third-party symbol.

.template.yaml

@@ -1,6 +1,6 @@
 # Autogenerated - DO NOT EDIT
 # Parameters of the project as generated from template
-_commit: 2555e1a
+_commit: fcdbd3c
 _src_path: gh:kraysent/python-template
 package_name: uploader
 project_name: uploader

AGENTS.md

@@ -79,3 +79,7 @@ The only exception to the above rule is for imports from the `typing` or `collec
 from typing import Any
 from collections.abc import Sequence
 ```
+
+### File changes
+
+When moving files as part of a refactor of any kind prefer using `git mv` whenver possible instead of removing + adding.

tests/test_formula_evaluate.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 
+import astropy.constants as const
 import astropy.units as u
 import numpy as np
 import pytest
@@ -124,6 +125,131 @@ def evaluate_expr(source: str, columns: dict[str, Col]) -> object:
         result_val=1.5,
         result_unit=u.dimensionless_unscaled,
     ),
+    EvalCase(
+        name="tan",
+        expression="tan(45 * deg)",
+        columns={},
+        result_val=1.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="asin",
+        expression='asin(col("float_col_dimless"))',
+        columns=_COLUMNS,
+        result_val=0.5236,
+        result_unit=u.rad,
+    ),
+    EvalCase(
+        name="acos",
+        expression="acos(1)",
+        columns={},
+        result_val=0.0,
+        result_unit=u.rad,
+    ),
+    EvalCase(
+        name="atan",
+        expression="atan(1)",
+        columns={},
+        result_val=0.7854,
+        result_unit=u.rad,
+    ),
+    EvalCase(
+        name="atan2",
+        expression="atan2(0, 1)",
+        columns={},
+        result_val=0.0,
+        result_unit=u.rad,
+    ),
+    EvalCase(
+        name="deg2rad",
+        expression="deg2rad(180)",
+        columns={},
+        result_val=np.pi,
+        result_unit=u.rad,
+    ),
+    EvalCase(
+        name="rad2deg",
+        expression="rad2deg(pi)",
+        columns={},
+        result_val=180.0,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="wrap360",
+        expression="wrap360(370 * deg)",
+        columns={},
+        result_val=10.0,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="wrap360_negative",
+        expression="wrap360(-10 * deg)",
+        columns={},
+        result_val=350.0,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="exp",
+        expression="exp(0)",
+        columns={},
+        result_val=1.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="log10",
+        expression="log10(100)",
+        columns={},
+        result_val=2.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="ln",
+        expression="ln(1)",
+        columns={},
+        result_val=0.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="pow",
+        expression="pow(2, 3)",
+        columns={},
+        result_val=8.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="max",
+        expression="max(2, 3)",
+        columns={},
+        result_val=3.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="min",
+        expression="min(2, 3)",
+        columns={},
+        result_val=2.0,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="max_equivalent_units",
+        expression="max(10 * deg, 1 * rad)",
+        columns={},
+        result_val=57.2958,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="vector_max",
+        expression='max(col("vec_col"), 2)',
+        columns=_COLUMNS,
+        result_val=[2.0, 2.0, 3.0],
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="error_max_incompatible_units",
+        expression="max(1 * deg, 1)",
+        columns={},
+        error=True,
+    ),
     EvalCase(
         name="vector_sqrt",
         expression='sqrt(col("vec_col"))',
@@ -138,6 +264,34 @@ def evaluate_expr(source: str, columns: dict[str, Col]) -> object:
         result_val=3.1416,
         result_unit=u.dimensionless_unscaled,
     ),
+    EvalCase(
+        name="speed_of_light",
+        expression="c",
+        columns={},
+        result_val=299792458.0,
+        result_unit=u.m / u.s,
+    ),
+    EvalCase(
+        name="solar_mass",
+        expression="M_sun",
+        columns={},
+        result_val=const.M_sun.value,
+        result_unit=u.kg,
+    ),
+    EvalCase(
+        name="parsec",
+        expression="pc",
+        columns={},
+        result_val=1.0,
+        result_unit=u.pc,
+    ),
+    EvalCase(
+        name="jansky",
+        expression="Jy",
+        columns={},
+        result_val=1.0,
+        result_unit=u.Jy,
+    ),
     EvalCase(name="error_missing_column_call", expression='col("missing")', columns={}, error=True),
     EvalCase(name="error_incompatible_units", expression="arcsec + mag", columns={}, error=True),
     EvalCase(
@@ -260,6 +414,59 @@ def evaluate_expr(source: str, columns: dict[str, Col]) -> object:
         result_val=1.5,
         result_unit=u.km / u.s,
     ),
+    EvalCase(
+        name="where_true",
+        expression='where(col("float_col") > 1, col("float_col"), 0)',
+        columns=_COLUMNS,
+        result_val=1.5,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="where_false",
+        expression='where(col("float_col") > 2, 10, col("float_col"))',
+        columns=_COLUMNS,
+        result_val=1.5,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="where_nested",
+        expression='where(col("float_col") > 2, 1, where(col("float_col") > 1, 2, 3))',
+        columns=_COLUMNS,
+        result_val=2,
+    ),
+    EvalCase(
+        name="where_angle",
+        expression='where(col("angle_col") > 180 * deg, col("angle_col") - 360 * deg, col("angle_col"))',
+        columns=_COLUMNS,
+        result_val=-170.0,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="where_string",
+        expression='where(col("string_col_2") == "M", col("string_col_1"), "x")',
+        columns=_COLUMNS,
+        result_val="NGC 123",
+    ),
+    EvalCase(
+        name="where_vector",
+        expression='where(col("vec_col") > 1.5, col("vec_col"), 0)',
+        columns=_COLUMNS,
+        result_val=[0.0, 2.0, 3.0],
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="where_vector_angle",
+        expression='where(col("vec_angle_col") > 90 * deg, col("vec_angle_col"), 0 * deg)',
+        columns=_COLUMNS,
+        result_val=[0.0, 0.0, 180.0],
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="where_vector_string",
+        expression='where(col("vec_string_a") == "IC", col("vec_string_b"), "x")',
+        columns=_COLUMNS,
+        result_val=["x", "456", "x"],
+    ),
     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),
 ]

tests/test_formula_parse.py

@@ -12,6 +12,8 @@
     ('"M " + col("id")', {"id"}),
     ("1 + 2", set()),
     ("pi + deg", set()),
+    ("M_sun * G / c ** 2", set()),
+    ("max(1, min(2, 3))", set()),
     ("", ExpressionSyntaxError),
     ("1 +", ExpressionSyntaxError),
     ("col(", ExpressionSyntaxError),
@@ -24,6 +26,9 @@
     ("e_logd25 + logd25", ExpressionSyntaxError),
     ("logd25 + logd25", ExpressionSyntaxError),
     ('3 * 10 ** col("logd25") * e_logd25 * arcsec', ExpressionSyntaxError),
+    ("atan2(1, 0) + wrap360(10 * deg)", set()),
+    ('where(col("a") > 0, col("a"), 0)', {"a"}),
+    ('where(col("a") == "M", col("b"), where(col("a") == "N", col("c"), "x"))', {"a", "b", "c"}),
     ("unknown(1)", ExpressionSyntaxError),
     ("col", ExpressionSyntaxError),
 ]

tests/test_formula_validate.py

@@ -9,7 +9,12 @@ 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("sqrt(4)") == []
+    assert validate_expression("atan2(1, 0) + pow(2, 3) + wrap360(10 * deg)") == []
     assert validate_expression('col("dist") * unit("Mpc")') == []
+    assert validate_expression("M_sun / L_sun * sigma") == []
+    assert validate_expression('where(col("v") > 0, col("v"), 0)') == []
+    assert validate_expression('where(col("a") == "M", 1, where(col("a") != "N", 2, 3))') == []
+    assert validate_expression("max(1, min(2, 3))") == []
     assert validate_expression("") == []
 
 

uploader/app/lib/formula/constants.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from typing import final
+
+import astropy.constants as const
+import astropy.units as u
+import numpy as np
+
+
+@final
+@dataclass(frozen=True)
+class ConstantDef:
+    name: str
+    value: u.Quantity
+    detail: str
+
+    @property
+    def insert(self) -> str:
+        return self.name
+
+
+NAMED_CONSTANTS: tuple[ConstantDef, ...] = (
+    ConstantDef("pi", np.pi * u.dimensionless_unscaled, "Pi"),
+    ConstantDef("c", const.c, "Speed of light"),
+    ConstantDef("G", const.G, "Gravitational constant"),
+    ConstantDef("h", const.h, "Planck constant"),
+    ConstantDef("k_B", const.k_B, "Boltzmann constant"),
+    ConstantDef("sigma", const.sigma_sb, "Stefan-Boltzmann constant"),
+    ConstantDef("m_e", const.m_e, "Electron mass"),
+    ConstantDef("m_p", const.m_p, "Proton mass"),
+    ConstantDef("au", 1 * u.au, "Astronomical unit"),
+    ConstantDef("pc", 1 * u.pc, "Parsec"),
+    ConstantDef("ly", 1 * u.lyr, "Light year"),
+    ConstantDef("eV", 1 * u.eV, "Electronvolt"),
+    ConstantDef("Jy", 1 * u.Jy, "Jansky"),
+    ConstantDef("M_sun", const.M_sun, "Solar mass"),
+    ConstantDef("R_sun", const.R_sun, "Nominal solar radius"),
+    ConstantDef("L_sun", const.L_sun, "Nominal solar luminosity"),
+    ConstantDef("M_earth", const.M_earth, "Earth mass"),
+    ConstantDef("M_jup", const.M_jup, "Jupiter mass"),
+    ConstantDef("deg", 1 * u.deg, "Degree"),
+    ConstantDef("rad", 1 * u.rad, "Radian"),
+    ConstantDef("arcmin", 1 * u.arcmin, "Arcminute"),
+    ConstantDef("arcsec", 1 * u.arcsec, "Arcsecond"),
+    ConstantDef("mag", 1 * u.mag, "Magnitude"),
+)