diff --git a/src/substrait/derivation_expression.py b/src/substrait/derivation_expression.py index 0d8b153f..be9d9690 100644 --- a/src/substrait/derivation_expression.py +++ b/src/substrait/derivation_expression.py @@ -2,10 +2,36 @@ from typing import Optional from antlr4 import CommonTokenStream, InputStream +from antlr4.error.ErrorListener import ErrorListener from substrait.type_pb2 import NamedStruct, Type from substrait_antlr.substrait_type.SubstraitTypeLexer import SubstraitTypeLexer from substrait_antlr.substrait_type.SubstraitTypeParser import SubstraitTypeParser + +class DerivationExpressionParseError(Exception): + """Raised when a derivation expression cannot be parsed.""" + + +class _RaisingErrorListener(ErrorListener): + """ANTLR error listener that raises instead of recovering. + + ANTLR's default behaviour on a syntax error is to print a warning to + stderr and continue with a partial (possibly ``None``-containing) parse + tree. We want the parse to fail loudly with a message naming the offending + input instead. + """ + + def __init__(self, expression: str): + super().__init__() + self._expression = expression + + def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e): + raise DerivationExpressionParseError( + f"Could not parse derivation expression {self._expression!r} " + f"at line {line}:{column}: {msg}" + ) from e + + # Binary operators keyed by the token text the grammar attaches to `op`. # (Integer division: type parameters like precision/scale are integers.) _BINARY_OPS = { @@ -245,9 +271,18 @@ def _evaluate(x, values: dict): def _parse(x: str): + error_listener = _RaisingErrorListener(x) + lexer = SubstraitTypeLexer(InputStream(x)) + lexer.removeErrorListeners() + lexer.addErrorListener(error_listener) + stream = CommonTokenStream(lexer) + parser = SubstraitTypeParser(stream) + parser.removeErrorListeners() + parser.addErrorListener(error_listener) + return parser.expr() diff --git a/tests/test_derivation_expression.py b/tests/test_derivation_expression.py index 680c7c2c..2caa62ea 100644 --- a/tests/test_derivation_expression.py +++ b/tests/test_derivation_expression.py @@ -1,6 +1,9 @@ +import re + +import pytest from substrait.type_pb2 import NamedStruct, Type -from substrait.derivation_expression import evaluate +from substrait.derivation_expression import DerivationExpressionParseError, evaluate def test_simple_arithmetic(): @@ -146,7 +149,7 @@ def test_struct_simple(): def test_nstruct_simple(): """Test named struct with field names and types.""" - result = evaluate("nStruct", {}) + result = evaluate("nStruct", {}) expected = NamedStruct( names=["a", "b"], struct=Type.Struct( @@ -160,9 +163,33 @@ def test_nstruct_simple(): assert result == expected +@pytest.mark.parametrize( + "expression", + [ + "decimal", + "interval_day", + "interval_compound", + "precision_time", + "precision_timestamp", + "precision_timestamp_tz", + ], +) +def test_parameterized_type_without_arguments_raises(expression): + """Precision-parameterized types used without their required arguments are + syntax errors and should raise a clear parse error, not fall through to the + confusing "Unknown token type " exception (see #216).""" + with pytest.raises(DerivationExpressionParseError, match=re.escape(expression)): + evaluate(expression) + + +def test_invalid_syntax_raises(): + with pytest.raises(DerivationExpressionParseError): + evaluate("i8 +") + + def test_nstruct_nested(): """Test named struct with nested struct field.""" - result = evaluate("nStruct>", {}) + result = evaluate("nStruct>", {}) expected = NamedStruct( names=["a", "b", "c"], struct=Type.Struct(