diff --git a/src/too/tools.py b/src/too/tools.py index 6f1f272..c7810be 100644 --- a/src/too/tools.py +++ b/src/too/tools.py @@ -20,6 +20,7 @@ import rich.progress from too.datamodel import too_dtypes +from too.exceptions import ValidationError if TYPE_CHECKING: @@ -60,10 +61,19 @@ def read_too_file( else: targets = path - if cast: - targets = targets.cast(too_dtypes) + try: + if cast: + targets = targets.cast(too_dtypes) + + targets = targets.sort("added_on", "too_id").select(list(too_dtypes)) + + except polars.exceptions.ColumnNotFoundError as err: + column = str(err).splitlines()[0] + raise ValidationError(f"Missing column {column!r} in input file.") + + except Exception as err: + raise ValidationError(f"Error reading input file: {err!r}") - targets = targets.sort("added_on", "too_id").select(list(too_dtypes)) return targets diff --git a/tests/test_tools.py b/tests/test_tools.py index f7567d8..8d9ef00 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -13,7 +13,9 @@ import polars import pytest -from too.tools import match_fields +from too.datamodel import too_dtypes +from too.exceptions import ValidationError +from too.tools import match_fields, read_too_file if TYPE_CHECKING: @@ -62,3 +64,22 @@ def test_match_fields_check_separation( # Most random coordinates will have some that do not fall inside our tiling. with pytest.raises(ValueError): match_fields(too_mock, database, "eta-6", check_separation=True) + + +def test_read_too_file_missing_column(): + df = polars.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + + with pytest.raises(ValidationError) as err: + read_too_file(df) + + assert "Missing column" in str(err) + + +def test_read_too_file_cast_failes(): + df = polars.DataFrame({name: None for name in too_dtypes}, schema=too_dtypes) + df = df.with_columns(g_mag=polars.Series(["hello"], dtype=polars.String)) + + with pytest.raises(ValidationError) as err: + read_too_file(df, cast=True) + + assert "Error reading input file" in str(err)