diff --git a/tests/test_designation_upload.py b/tests/test_designation_upload.py new file mode 100644 index 0000000..659677b --- /dev/null +++ b/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"] diff --git a/tests/test_formula_evaluate.py b/tests/test_formula_evaluate.py new file mode 100644 index 0000000..8c438ba --- /dev/null +++ b/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) diff --git a/tests/test_formula_parse.py b/tests/test_formula_parse.py new file mode 100644 index 0000000..7a0ccb0 --- /dev/null +++ b/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 diff --git a/uploader/app/lib/formula/__init__.py b/uploader/app/lib/formula/__init__.py new file mode 100644 index 0000000..b5f3f85 --- /dev/null +++ b/uploader/app/lib/formula/__init__.py @@ -0,0 +1,20 @@ +from uploader.app.lib.formula.core import Expression, evaluate, parse +from uploader.app.lib.formula.errors import ( + ExpressionError, + ExpressionEvaluationError, + ExpressionSyntaxError, +) +from uploader.app.lib.formula.namespace import expression_syntax_help +from uploader.app.lib.formula.values import Value, column_quantity + +__all__ = [ + "Expression", + "ExpressionError", + "ExpressionEvaluationError", + "ExpressionSyntaxError", + "Value", + "column_quantity", + "evaluate", + "expression_syntax_help", + "parse", +] diff --git a/uploader/app/lib/formula/core.py b/uploader/app/lib/formula/core.py new file mode 100644 index 0000000..cf1fdfa --- /dev/null +++ b/uploader/app/lib/formula/core.py @@ -0,0 +1,72 @@ +import ast +from collections.abc import Mapping +from dataclasses import dataclass +from types import CodeType +from typing import final + +import astropy.units as u + +from uploader.app.lib.formula.errors import ExpressionEvaluationError, ExpressionSyntaxError +from uploader.app.lib.formula.namespace import COL_FUNCTION, FUNCTIONS, NAMED_CONSTANTS, build_namespace +from uploader.app.lib.formula.values import Value + + +def _column_from_call(node: ast.Call) -> str | None: + if not isinstance(node.func, ast.Name) or node.func.id != COL_FUNCTION: + return None + if node.keywords or len(node.args) != 1: + raise ValueError(f"{COL_FUNCTION}() takes exactly one string argument") + arg = node.args[0] + if not isinstance(arg, ast.Constant) or not isinstance(arg.value, str): + raise ValueError(f"{COL_FUNCTION}() argument must be a string literal") + return arg.value + + +@final +@dataclass(frozen=True) +class Expression: + referenced_columns: frozenset[str] + code: CodeType + + +def parse(source: str) -> Expression: + try: + tree = ast.parse(source.strip(), mode="eval") + except SyntaxError as e: + raise ExpressionSyntaxError(str(e)) from e + try: + referenced_columns = frozenset(_ColumnCollector().collect(tree.body)) + except ValueError as e: + raise ExpressionSyntaxError(str(e)) from e + code = compile(tree, "", "eval") + return Expression(referenced_columns=referenced_columns, code=code) + + +def evaluate(expression: Expression, columns: Mapping[str, Value]) -> Value: + try: + return eval(expression.code, build_namespace(columns)) # noqa: S307 + except (KeyError, NameError, TypeError, ValueError, ZeroDivisionError, u.UnitsError) as e: + raise ExpressionEvaluationError(str(e)) from e + + +@final +class _ColumnCollector(ast.NodeVisitor): + def __init__(self) -> None: + self.columns: set[str] = set() + + def collect(self, node: ast.AST) -> set[str]: + self.visit(node) + return self.columns + + def visit_Call(self, node: ast.Call) -> None: + column = _column_from_call(node) + if column is not None: + self.columns.add(column) + return + for arg in node.args: + self.visit(arg) + + def visit_Name(self, node: ast.Name) -> None: + if node.id in NAMED_CONSTANTS or node.id in FUNCTIONS: + return + self.columns.add(node.id) diff --git a/uploader/app/lib/formula/errors.py b/uploader/app/lib/formula/errors.py new file mode 100644 index 0000000..aee0e07 --- /dev/null +++ b/uploader/app/lib/formula/errors.py @@ -0,0 +1,10 @@ +class ExpressionError(Exception): + pass + + +class ExpressionSyntaxError(ExpressionError): + pass + + +class ExpressionEvaluationError(ExpressionError): + pass diff --git a/uploader/app/lib/formula/namespace.py b/uploader/app/lib/formula/namespace.py new file mode 100644 index 0000000..cc73c2a --- /dev/null +++ b/uploader/app/lib/formula/namespace.py @@ -0,0 +1,71 @@ +from collections.abc import Mapping + +import astropy.constants as const +import astropy.units as u +import numpy as np + +from uploader.app.lib.formula.values import Value + +COL_FUNCTION = "col" + +NAMED_CONSTANTS: dict[str, u.Quantity] = { + "pi": np.pi * u.dimensionless_unscaled, + "c": const.c, + "deg": 1 * u.deg, + "rad": 1 * u.rad, + "arcmin": 1 * u.arcmin, + "arcsec": 1 * u.arcsec, + "mag": 1 * u.mag, +} + + +def _scalar_to_str(value: float | int | np.number) -> str: + numeric = float(value) + if numeric.is_integer(): + return str(int(numeric)) + return str(numeric) + + +def _formula_str(value: Value) -> str | np.ndarray: + if isinstance(value, str): + return value + if isinstance(value, u.Quantity): + scalar = value.value + if isinstance(scalar, np.ndarray): + return np.asarray([_scalar_to_str(x) for x in scalar]) + return _scalar_to_str(scalar) + return np.asarray([_scalar_to_str(x) for x in value]) + + +FUNCTIONS: dict[str, object] = { + "sin": np.sin, + "cos": np.cos, + "str": _formula_str, +} + + +def build_namespace(columns: Mapping[str, Value]) -> dict[str, object]: + bare = {name: value for name, value in columns.items() if name.isidentifier()} + return { + "__builtins__": {}, + COL_FUNCTION: lambda name: columns[name], + **bare, + **NAMED_CONSTANTS, + **FUNCTIONS, + } + + +def expression_syntax_help() -> str: + constants = ", ".join(sorted(NAMED_CONSTANTS)) + return ( + f'Use {COL_FUNCTION}("name") or bare identifiers to refer to rawdata columns ' + '(e.g. col("a"), e_logd25).\n' + "Bare identifiers that match predefined constants use those constants.\n" + "Operators: + - * / ** %.\n" + "Functions: sin(x), cos(x) (argument must be an angle), str(x).\n" + "Numbers are dimensionless.\n" + "String literals and + concatenation are supported.\n" + 'Modulo divisors must carry units (e.g. col("pa") % (180 * deg)).\n' + "Log columns (mag/dex) yield the bare exponent; multiply by the scale yourself.\n" + f"Available constants: {constants}." + ) diff --git a/uploader/app/lib/formula/values.py b/uploader/app/lib/formula/values.py new file mode 100644 index 0000000..102a4ba --- /dev/null +++ b/uploader/app/lib/formula/values.py @@ -0,0 +1,26 @@ +from collections.abc import Sequence + +import astropy.units as u +import numpy as np +from astropy.units.function.core import FunctionUnitBase + +type Value = u.Quantity | str | np.ndarray + + +def _is_logarithmic_column_unit(unit: u.Unit) -> bool: + return unit == u.mag or unit == u.dex or isinstance(unit, FunctionUnitBase) + + +def column_quantity(value: float | str | Sequence[float] | Sequence[str], unit: str) -> Value: + if isinstance(value, str): + return value + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + if all(isinstance(x, str) for x in value): + return np.asarray(value) + numeric = np.asarray(value, dtype=float) + if not unit: + return numeric * u.dimensionless_unscaled + parsed = u.Unit(unit) + if _is_logarithmic_column_unit(parsed): + return numeric * u.dimensionless_unscaled + return numeric * parsed diff --git a/uploader/app/structured/designations/upload.py b/uploader/app/structured/designations/upload.py index 34a55a4..0c99ee7 100644 --- a/uploader/app/structured/designations/upload.py +++ b/uploader/app/structured/designations/upload.py @@ -1,18 +1,31 @@ +import csv +import pathlib from collections.abc import Callable +from typing import Any +import astropy.units as u import matplotlib.pyplot as plt +import numpy as np from psycopg import sql import uploader.app.action_description as action_description import uploader.app.report as report from uploader.app import log from uploader.app.display import format_table +from uploader.app.lib.formula import ( + Expression, + ExpressionEvaluationError, + Value, + column_quantity, + evaluate, + parse, +) from uploader.app.lib.rawdata import rawdata_batches from uploader.app.storage import PgStorage from uploader.app.structured.designations.rules import RULES, match from uploader.app.upload import handle_call from uploader.clients.gen.client import adminapi -from uploader.clients.gen.client.adminapi.api.default import save_structured_data +from uploader.clients.gen.client.adminapi.api.default import get_table, save_structured_data from uploader.clients.gen.client.adminapi.models.save_structured_data_request import ( SaveStructuredDataRequest, ) @@ -114,17 +127,96 @@ def pct(n: int) -> float: report_func(report.DoneEvent(message=summary)) +def _fetch_column_units( + client: adminapi.AuthenticatedClient, + table_name: str, +) -> tuple[set[str], dict[str, str]]: + resp = handle_call(get_table.sync_detailed(client=client, table_name=table_name)) + column_names: set[str] = set() + column_units: dict[str, str] = {} + for col in resp.data.column_info: + column_names.add(col.name) + if isinstance(col.unit, str): + column_units[col.name] = col.unit + return column_names, column_units + + +def _validate_columns( + table_name: str, + needed_cols: set[str], + column_names: set[str], +) -> None: + missing = sorted(col for col in needed_cols if col not in column_names) + if missing: + raise RuntimeError(f"Table {table_name} has no column(s): {missing}") + + +def _build_column_values( + row: dict[str, Any], + referenced_columns: frozenset[str], + column_units: dict[str, str], +) -> dict[str, Value]: + return {col: column_quantity(row[col], column_units.get(col, "")) for col in referenced_columns} + + +def _designation_string(value: Value) -> str: + if isinstance(value, str): + return value.strip() + if isinstance(value, u.Quantity): + scalar = value.value + if isinstance(scalar, np.ndarray) and scalar.shape != (): + raise RuntimeError("designation expression must evaluate to a scalar value per row") + return str(scalar).strip() + raise RuntimeError("designation expression must evaluate to a scalar value per row") + + +def _evaluate_designation( + parsed: Expression, + row: dict[str, Any], + column_units: dict[str, str], +) -> str: + columns = _build_column_values(row, parsed.referenced_columns, column_units) + try: + return _designation_string(evaluate(parsed, columns)) + except ExpressionEvaluationError as e: + raise RuntimeError( + f"failed to evaluate expression for row {row['hyperleda_internal_id']}: {e}", + ) from e + + +class _DesignationOutputWriter: + def __init__(self, path: str) -> None: + output_path = pathlib.Path(path) + output_path.parent.mkdir(parents=True, exist_ok=True) + self._file = output_path.open("w", newline="") + self._writer = csv.writer(self._file) + self._writer.writerow(["id", "designation"]) + + def write_batch(self, ids: list[str], designations: list[list[str]]) -> None: + for record_id, designation in zip(ids, designations, strict=True): + self._writer.writerow([record_id, designation[0]]) + + def close(self) -> None: + self._file.close() + + def upload_designations( storage: PgStorage, table_name: str, - column_name: str, + expression: str, batch_size: int, client: adminapi.AuthenticatedClient, *, write: bool = False, print_unmatched: bool = False, + output_file: str = "", report_func: Callable[[report.Event], None], ) -> int: + parsed = parse(expression) + needed_cols = set(parsed.referenced_columns) + column_names, column_units = _fetch_column_units(client, table_name) + _validate_columns(table_name, needed_cols, column_names) + rule_counts: dict[str, int] = {r.name: 0 for r in RULES} unmatched = 0 total_count = 0 @@ -135,69 +227,79 @@ def upload_designations( total_count = int(cnt[0]["cnt"]) if cnt else 0 processed_rows = 0 + output_writer = _DesignationOutputWriter(output_file) if output_file else None + + try: + for rows in rawdata_batches(storage, table_name, sorted(needed_cols), batch_size): + batch_ids: list[str] = [] + batch_names: list[list[str]] = [] - for rows in rawdata_batches(storage, table_name, [column_name], batch_size): - batch_ids: list[str] = [] - batch_names: list[list[str]] = [] - - for row in rows: - internal_id = row["hyperleda_internal_id"] - name_val = row[column_name] - if name_val is None or (isinstance(name_val, str) and not name_val.strip()): - unmatched += 1 - continue - name_str = str(name_val).strip() - match_result = match(name_str) - if match_result is not None: - transformed, rule_name = match_result - rule_counts[rule_name] += 1 - else: - unmatched += 1 - transformed = name_str - if print_unmatched: - report_func(report.LogEvent(message=name_str)) - batch_ids.append(internal_id) - batch_names.append([transformed]) - - if write and batch_ids: - handle_call( - save_structured_data.sync_detailed( - client=client, - body=action_description.apply( - SaveStructuredDataRequest( - catalog="designation", - columns=["design"], - ids=batch_ids, - data=batch_names, + for row in rows: + internal_id = row["hyperleda_internal_id"] + if any(row[col] is None for col in needed_cols): + unmatched += 1 + continue + name_str = _evaluate_designation(parsed, row, column_units) + if not name_str: + unmatched += 1 + continue + match_result = match(name_str) + if match_result is not None: + transformed, rule_name = match_result + rule_counts[rule_name] += 1 + else: + unmatched += 1 + transformed = name_str + if print_unmatched: + report_func(report.LogEvent(message=name_str)) + batch_ids.append(internal_id) + batch_names.append([transformed]) + + if write and batch_ids: + handle_call( + save_structured_data.sync_detailed( + client=client, + body=action_description.apply( + SaveStructuredDataRequest( + catalog="designation", + columns=["design"], + ids=batch_ids, + data=batch_names, + ), ), - ), + ) ) - ) - processed_rows += len(rows) - total_so_far = sum(rule_counts.values()) + unmatched - - def total_pct(n: int, t: int = total_so_far) -> float: - return (100.0 * n / t) if t else 0.0 - - log.logger.info( - "processed batch", - total=total_so_far, - matched=sum(rule_counts.values()), - matched_pct=round(total_pct(sum(rule_counts.values())), 1), - unmatched=unmatched, - unmatched_pct=round(total_pct(unmatched), 1), - ) - progress_pct = int(100 * processed_rows / total_count) if total_count else 0 - _report_batch_progress( - report_func, - rows_read=len(rows), - total_so_far=total_so_far, - matched=sum(rule_counts.values()), - unmatched=unmatched, - progress_pct=progress_pct, - rule_counts=rule_counts, - ) + if output_writer is not None and batch_ids: + output_writer.write_batch(batch_ids, batch_names) + + processed_rows += len(rows) + total_so_far = sum(rule_counts.values()) + unmatched + + def total_pct(n: int, t: int = total_so_far) -> float: + return (100.0 * n / t) if t else 0.0 + + log.logger.info( + "processed batch", + total=total_so_far, + matched=sum(rule_counts.values()), + matched_pct=round(total_pct(sum(rule_counts.values())), 1), + unmatched=unmatched, + unmatched_pct=round(total_pct(unmatched), 1), + ) + progress_pct = int(100 * processed_rows / total_count) if total_count else 0 + _report_batch_progress( + report_func, + rows_read=len(rows), + total_so_far=total_so_far, + matched=sum(rule_counts.values()), + unmatched=unmatched, + progress_pct=progress_pct, + rule_counts=rule_counts, + ) + finally: + if output_writer is not None: + output_writer.close() total = sum(rule_counts.values()) + unmatched _report_rule_distribution(report_func, rule_counts, unmatched, total) diff --git a/uploader/forms/structured_designation.py b/uploader/forms/structured_designation.py index 632ce13..1d34d0e 100644 --- a/uploader/forms/structured_designation.py +++ b/uploader/forms/structured_designation.py @@ -21,14 +21,22 @@ class StructuredDesignationAdvancedSettings(BaseModel): title="Log unmatched names", description="Append each unmatched name to the log stream.", ) + output_file: str = Field( + default="", + title="Output file", + description="If set, write id and designation pairs to this file path.", + ) class StructuredDesignationForm(BaseModel): table_name: str = Field(..., title="Name of the table") - column_name: str = Field( + expression: str = Field( ..., - title="Object name column", - description="Name of the column that represents object designation in the table.", + title="Designation expression", + description=( + "Expression yielding the object designation per row. " + 'Examples: designation, col("weird name"), prefix + " " + number.' + ), ) write: bool = Field( default=False, @@ -61,10 +69,11 @@ def handle_structured_designation( run_upload_designations( storage, f.table_name.strip(), - f.column_name.strip(), + f.expression.strip(), advanced.batch_size, client, write=f.write, print_unmatched=advanced.print_unmatched, + output_file=advanced.output_file.strip(), report_func=report_func, ) diff --git a/uploader/task_registry.py b/uploader/task_registry.py index fc6670a..c4fdba7 100644 --- a/uploader/task_registry.py +++ b/uploader/task_registry.py @@ -1,4 +1,4 @@ -from uploader.app.lib.expression import expression_syntax_help +from uploader.app.lib.formula import expression_syntax_help from uploader.forms.authenticate import AuthenticateForm, handle_authenticate from uploader.forms.crossmatch_layered import CrossmatchLayeredForm, handle_crossmatch_layered from uploader.forms.structured_designation import ( @@ -80,7 +80,9 @@ def register_all_tasks() -> None: TaskDefinition( id="upload-structured-designation", title="Designations", - description="Convert designations to common format and upload them to the database.", + description=( + f"Convert designations to common format and upload them to the database.\n\n{expression_syntax_help()}" + ), form_model=StructuredDesignationForm, handler=handle_structured_designation, group="Catalogs",