add ability to make designation expressions

tests/test_designation_upload.py

@@ -0,0 +1,151 @@
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+import pytest
+
+from uploader.app.structured.designations.upload import upload_designations
+from uploader.clients.gen.client import adminapi
+
+
+def _mock_storage(total: int = 1) -> Mock:
+    storage = Mock()
+    storage.query.return_value = [{"cnt": total}]
+    return storage
+
+
+def _mock_client() -> Mock:
+    return Mock(spec=adminapi.AuthenticatedClient)
+
+
+@patch("uploader.app.structured.designations.upload.rawdata_batches")
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_simple_column_expression(
+    mock_fetch_column_units: Mock,
+    mock_rawdata_batches: Mock,
+) -> None:
+    mock_fetch_column_units.return_value = ({"name"}, {"name": ""})
+    mock_rawdata_batches.return_value = iter(
+        [[{"hyperleda_internal_id": "1", "name": "NGC 123"}]],
+    )
+
+    total = upload_designations(
+        _mock_storage(),
+        "test_table",
+        "name",
+        100,
+        _mock_client(),
+        report_func=lambda _: None,
+    )
+
+    assert total == 1
+    mock_rawdata_batches.assert_called_once()
+    assert mock_rawdata_batches.call_args.args[2] == ["name"]
+
+
+@patch("uploader.app.structured.designations.upload.rawdata_batches")
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_composed_string_expression(
+    mock_fetch_column_units: Mock,
+    mock_rawdata_batches: Mock,
+) -> None:
+    mock_fetch_column_units.return_value = ({"prefix", "number"}, {"prefix": "", "number": ""})
+    mock_rawdata_batches.return_value = iter(
+        [[{"hyperleda_internal_id": "1", "prefix": "NGC", "number": "123"}]],
+    )
+
+    total = upload_designations(
+        _mock_storage(),
+        "test_table",
+        'prefix + " " + number',
+        100,
+        _mock_client(),
+        report_func=lambda _: None,
+    )
+
+    assert total == 1
+
+
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_missing_referenced_columns(mock_fetch_column_units: Mock) -> None:
+    mock_fetch_column_units.return_value = ({"other"}, {})
+
+    with pytest.raises(RuntimeError, match="has no column\\(s\\): \\['name'\\]"):
+        upload_designations(
+            _mock_storage(),
+            "test_table",
+            "name",
+            100,
+            _mock_client(),
+            report_func=lambda _: None,
+        )
+
+
+@patch("uploader.app.structured.designations.upload.rawdata_batches")
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_null_referenced_values_counted_as_unmatched(
+    mock_fetch_column_units: Mock,
+    mock_rawdata_batches: Mock,
+) -> None:
+    mock_fetch_column_units.return_value = ({"name"}, {"name": ""})
+    mock_rawdata_batches.return_value = iter(
+        [[{"hyperleda_internal_id": "1", "name": None}]],
+    )
+
+    total = upload_designations(
+        _mock_storage(),
+        "test_table",
+        "name",
+        100,
+        _mock_client(),
+        report_func=lambda _: None,
+    )
+
+    assert total == 1
+
+
+@patch("uploader.app.structured.designations.upload.rawdata_batches")
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_expression_evaluation_error_becomes_runtime_error(
+    mock_fetch_column_units: Mock,
+    mock_rawdata_batches: Mock,
+) -> None:
+    mock_fetch_column_units.return_value = ({"text_col", "num_col"}, {"text_col": "", "num_col": ""})
+    mock_rawdata_batches.return_value = iter(
+        [[{"hyperleda_internal_id": "1", "text_col": "NGC 123", "num_col": 1.5}]],
+    )
+
+    with pytest.raises(RuntimeError, match="failed to evaluate expression for row 1"):
+        upload_designations(
+            _mock_storage(),
+            "test_table",
+            "text_col + num_col",
+            100,
+            _mock_client(),
+            report_func=lambda _: None,
+        )
+
+
+@patch("uploader.app.structured.designations.upload.rawdata_batches")
+@patch("uploader.app.structured.designations.upload._fetch_column_units")
+def test_output_file_writes_id_designation_pairs(
+    mock_fetch_column_units: Mock,
+    mock_rawdata_batches: Mock,
+    tmp_path: Path,
+) -> None:
+    mock_fetch_column_units.return_value = ({"name"}, {"name": ""})
+    mock_rawdata_batches.return_value = iter(
+        [[{"hyperleda_internal_id": "id-1", "name": "NGC 123"}]],
+    )
+    output_file = tmp_path / "designations.csv"
+
+    upload_designations(
+        _mock_storage(),
+        "test_table",
+        "name",
+        100,
+        _mock_client(),
+        output_file=str(output_file),
+        report_func=lambda _: None,
+    )
+
+    assert output_file.read_text().splitlines() == ["id,designation", "id-1,NGC 123"]

tests/test_formula_evaluate.py

@@ -0,0 +1,211 @@
+from dataclasses import dataclass
+
+import astropy.units as u
+import numpy as np
+import pytest
+
+from uploader.app.lib.formula import ExpressionEvaluationError, column_quantity, evaluate, parse
+
+
+@dataclass
+class Col:
+    value: float | str | list[float] | list[str]
+    unit: str = ""
+
+
+@dataclass
+class EvalCase:
+    expression: str
+    columns: dict[str, Col]
+    result_val: str | float | list[float] | list[str] | None = None
+    result_unit: u.Unit | None = None
+    error: bool = False
+    name: str = ""
+
+
+def evaluate_expr(source: str, columns: dict[str, Col]) -> object:
+    built = {name: column_quantity(col.value, col.unit) for name, col in columns.items()}
+    return evaluate(parse(source), built)
+
+
+_COLUMNS: dict[str, Col] = {
+    "float_col": Col(1.5),
+    "angle_col": Col(190.0, "deg"),
+    "float_col_mag": Col(1.5, "mag"),
+    "float_col_dex": Col(0.697, "dex(0.1 arcmin)"),
+    "brightness_col": Col(23.162, "mag / arcsec2"),
+    "string_col_1": Col("NGC 123"),
+    "string_col_2": Col("M"),
+    "float_col_dimless": Col(0.5),
+    "vec_col": Col([1.0, 2.0, 3.0]),
+    "vec_angle_col": Col([0.0, 90.0, 180.0], "deg"),
+    "vec_string_a": Col(["NGC", "IC", "M"]),
+    "vec_string_b": Col(["123", "456", "789"]),
+}
+
+
+EVAL_CASES: list[EvalCase] = [
+    EvalCase(
+        name="dimentionless_expr",
+        expression='3 * 10 ** col("float_col") * arcsec',
+        columns=_COLUMNS,
+        result_val=94.8683,
+        result_unit=u.arcsec,
+    ),
+    EvalCase(
+        name="modulo_edgecase",
+        expression='col("angle_col") % (180.0 * deg)',
+        columns=_COLUMNS,
+        result_val=10.0,
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="single_column",
+        expression='col("brightness_col")',
+        columns=_COLUMNS,
+        result_val=23.162,
+        result_unit=u.Unit("mag/arcsec2"),
+    ),
+    EvalCase(name="string_literal", expression='"abc"', columns={}, result_val="abc"),
+    EvalCase(
+        name="string_column",
+        expression='col("string_col_1")',
+        columns=_COLUMNS,
+        result_val="NGC 123",
+    ),
+    EvalCase(
+        name="string_concat_columns",
+        expression='col("string_col_2") + " " + col("string_col_1")',
+        columns=_COLUMNS,
+        result_val="M NGC 123",
+    ),
+    EvalCase(
+        name="string_concat_with_str",
+        expression='col("string_col_2") + " " + str(col("float_col"))',
+        columns=_COLUMNS,
+        result_val="M 1.5",
+    ),
+    EvalCase(
+        name="str_on_string_column",
+        expression='str(col("string_col_1"))',
+        columns=_COLUMNS,
+        result_val="NGC 123",
+    ),
+    EvalCase(
+        name="str_integer",
+        expression='str(col("num"))',
+        columns={"num": Col(495444.0)},
+        result_val="495444",
+    ),
+    EvalCase(
+        name="str_keeps_fraction",
+        expression='str(col("float_col"))',
+        columns=_COLUMNS,
+        result_val="1.5",
+    ),
+    EvalCase(
+        name="function_unit",
+        expression='col("float_col_dex")',
+        columns=_COLUMNS,
+        result_val=0.697,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="nested_functions",
+        expression='sin(col("angle_col"))',
+        columns=_COLUMNS,
+        result_val=-0.1736,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="constant",
+        expression="pi",
+        columns={},
+        result_val=3.1416,
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(name="error_missing_column_call", expression='col("missing")', columns={}, error=True),
+    EvalCase(name="error_incompatible_units", expression="arcsec + mag", columns={}, error=True),
+    EvalCase(
+        name="error_string_plus_number",
+        expression='col("string_col_1") + col("float_col")',
+        columns=_COLUMNS,
+        error=True,
+    ),
+    EvalCase(
+        name="error_modulo_dimensionless_divisor",
+        expression='col("angle_col") % 180.0',
+        columns=_COLUMNS,
+        error=True,
+    ),
+    EvalCase(
+        name="error_trig_on_dimensionless",
+        expression='sin(col("float_col_dimless"))',
+        columns=_COLUMNS,
+        error=True,
+    ),
+    EvalCase(
+        name="vector_column_with_unit",
+        expression='col("vec_angle_col")',
+        columns=_COLUMNS,
+        result_val=[0.0, 90.0, 180.0],
+        result_unit=u.deg,
+    ),
+    EvalCase(
+        name="vector_arithmetic",
+        expression='3 * 10 ** col("vec_col") * arcsec',
+        columns=_COLUMNS,
+        result_val=[30.0, 300.0, 3000.0],
+        result_unit=u.arcsec,
+    ),
+    EvalCase(
+        name="vector_trig",
+        expression='sin(col("vec_angle_col"))',
+        columns=_COLUMNS,
+        result_val=[0.0, 1.0, 0.0],
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="scalar_vector_broadcast",
+        expression='col("float_col") + col("vec_col")',
+        columns=_COLUMNS,
+        result_val=[2.5, 3.5, 4.5],
+        result_unit=u.dimensionless_unscaled,
+    ),
+    EvalCase(
+        name="vector_string_concat",
+        expression='col("vec_string_a") + " " + col("vec_string_b")',
+        columns=_COLUMNS,
+        result_val=["NGC 123", "IC 456", "M 789"],
+    ),
+]
+
+
+@pytest.mark.parametrize("case", EVAL_CASES, ids=[case.name for case in EVAL_CASES])
+def test_evaluate(case: EvalCase) -> None:
+    if case.error:
+        with pytest.raises(ExpressionEvaluationError):
+            evaluate_expr(case.expression, case.columns)
+        return
+
+    result = evaluate_expr(case.expression, case.columns)
+
+    if case.result_unit is None:
+        if isinstance(case.result_val, list):
+            assert isinstance(result, np.ndarray)
+            np.testing.assert_array_equal(result, case.result_val)
+        else:
+            assert result == case.result_val
+        return
+
+    assert isinstance(result, u.Quantity)
+    assert result.unit == case.result_unit
+    if isinstance(case.result_val, list):
+        np.testing.assert_allclose(
+            np.asarray(result.value),
+            np.asarray(case.result_val),
+            rtol=1e-4,
+            atol=1e-10,
+        )
+    else:
+        np.testing.assert_almost_equal(result.value, case.result_val, decimal=4)

tests/test_formula_parse.py

@@ -0,0 +1,32 @@
+import pytest
+
+from uploader.app.lib.formula import ExpressionSyntaxError, parse
+
+PARSE_CASES: list[tuple[str, set[str] | type[ExpressionSyntaxError]]] = [
+    ("e_logd25 + logd25", {"e_logd25", "logd25"}),
+    ('col("a") + col("b")', {"a", "b"}),
+    ('col("weird name")', {"weird name"}),
+    ('sin(col("pa")) + pi', {"pa"}),
+    ("logd25 + logd25", {"logd25"}),
+    ('3 * 10 ** col("logd25") * e_logd25 * arcsec', {"logd25", "e_logd25"}),
+    ('"M " + col("id")', {"id"}),
+    ("1 + 2", set()),
+    ("", ExpressionSyntaxError),
+    ("1 +", ExpressionSyntaxError),
+    ("col(", ExpressionSyntaxError),
+    ("* 2", ExpressionSyntaxError),
+    ("a = 1", ExpressionSyntaxError),
+    ("col()", ExpressionSyntaxError),
+    ("col(x)", ExpressionSyntaxError),
+    ('col("a", "b")', ExpressionSyntaxError),
+    ("col(1)", ExpressionSyntaxError),
+]
+
+
+@pytest.mark.parametrize("source,expected", PARSE_CASES)
+def test_parse(source: str, expected: set[str] | type[ExpressionSyntaxError]) -> None:
+    if isinstance(expected, type):
+        with pytest.raises(expected):
+            parse(source)
+    else:
+        assert parse(source).referenced_columns == expected