fix: mypy errors and tests

Author: jaimeferj

pyiceberg/expressions/__init__.py

@@ -30,10 +30,11 @@
     Type,
     TypeVar,
     Union,
+    cast,
 )
 from typing import Literal as TypingLiteral
 
-from pydantic import ConfigDict, Field, field_serializer, field_validator
+from pydantic import ConfigDict, Field, field_validator
 
 from pyiceberg.expressions.literals import (
     AboveMax,
@@ -751,31 +752,50 @@ def as_bound(self) -> Type[BoundNotIn[L]]:
 
 class LiteralPredicate(IcebergBaseModel, UnboundPredicate[L], ABC):
     type: TypingLiteral["lt", "lt-eq", "gt", "gt-eq", "eq", "not-eq", "starts-with", "not-starts-with"] = Field(alias="type")
-    term: UnboundTerm[L]
-    literal: Literal[L] = Field(serialization_alias="value")
+    term: UnboundTerm[Any]
+    value: Literal[L] = Field(alias="literal", serialization_alias="value")
 
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+    model_config = ConfigDict(populate_by_name=True, frozen=True, arbitrary_types_allowed=True)
+
+    def __init__(
+        self,
+        term: Union[str, UnboundTerm[Any], BoundReference[Any]],
+        literal: Union[L, Literal[L], None] = None,
+        **data: Any,
+    ) -> None:  # pylint: disable=W0621
+        extra = dict(data)
+
+        literal_candidates = []
+        if literal is not None:
+            literal_candidates.append(literal)
+        if "literal" in extra:
+            literal_candidates.append(extra.pop("literal"))
+        if "value" in extra:
+            literal_candidates.append(extra.pop("value"))
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        if args:
-            if len(args) != 2:
-                raise TypeError("Expected (term, literal)")
-            kwargs = {"term": args[0], "literal": args[1], **kwargs}
-        super().__init__(**kwargs)
+        literal_candidates = [candidate for candidate in literal_candidates if candidate is not None]
+
+        if not literal_candidates:
+            raise TypeError("LiteralPredicate requires a literal or value argument")
+        if len(literal_candidates) > 1:
+            raise TypeError("literal/value provided multiple times")
+
+        init = cast("Callable[..., None]", IcebergBaseModel.__init__)
+        init(self, term=_to_unbound_term(term), literal=_to_literal(literal_candidates[0]), **extra)
 
     @field_validator("term", mode="before")
     @classmethod
-    def _coerce_term(cls, v: Any) -> UnboundTerm[Any]:
-        return _to_unbound_term(v)
+    def _convert_term(cls, value: Any) -> UnboundTerm[Any]:
+        return _to_unbound_term(value)
 
-    @field_validator("literal", mode="before")
+    @field_validator("value", mode="before")
     @classmethod
-    def _coerce_literal(cls, v: Union[L, Literal[L]]) -> Literal[L]:
-        return _to_literal(v)
+    def _convert_value(cls, value: Any) -> Literal[Any]:
+        return _to_literal(value)
 
-    @field_serializer("literal")
-    def ser_literal(self, literal: Literal[L]) -> str:
-        return "Any"
+    @property
+    def literal(self) -> Literal[L]:
+        return self.value
 
     def bind(self, schema: Schema, case_sensitive: bool = True) -> BoundLiteralPredicate[L]:
         bound_term = self.term.bind(schema, case_sensitive)

pyiceberg/transforms.py

@@ -120,7 +120,7 @@ def _try_import(module_name: str, extras_name: Optional[str] = None) -> types.Mo
         raise NotInstalledError(msg) from None
 
 
-def _transform_literal(func: Callable[[L], L], lit: Literal[L]) -> Literal[L]:
+def _transform_literal(func: Callable[[Any], Any], lit: Literal[L]) -> Literal[L]:
     """Small helper to upwrap the value from the literal, and wrap it again."""
     return literal(func(lit.value))
 

tests/expressions/test_evaluator.py

@@ -22,6 +22,7 @@
 from pyiceberg.conversions import to_bytes
 from pyiceberg.expressions import (
     And,
+    BooleanExpression,
     EqualTo,
     GreaterThan,
     GreaterThanOrEqual,
@@ -30,6 +31,7 @@
     IsNull,
     LessThan,
     LessThanOrEqual,
+    LiteralPredicate,
     Not,
     NotEqualTo,
     NotIn,
@@ -301,7 +303,7 @@ def test_missing_stats() -> None:
         upper_bounds=None,
     )
 
-    expressions = [
+    expressions: list[BooleanExpression] = [
         LessThan("no_stats", 5),
         LessThanOrEqual("no_stats", 30),
         EqualTo("no_stats", 70),
@@ -324,7 +326,7 @@ def test_zero_record_file_stats(schema_data_file: Schema) -> None:
         file_path="file_1.parquet", file_format=FileFormat.PARQUET, partition=Record(), record_count=0
     )
 
-    expressions = [
+    expressions: list[BooleanExpression] = [
         LessThan("no_stats", 5),
         LessThanOrEqual("no_stats", 30),
         EqualTo("no_stats", 70),
@@ -683,26 +685,27 @@ def data_file_nan() -> DataFile:
 
 
 def test_inclusive_metrics_evaluator_less_than_and_less_than_equal(schema_data_file_nan: Schema, data_file_nan: DataFile) -> None:
-    for operator in [LessThan, LessThanOrEqual]:
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+    operators: tuple[type[LiteralPredicate[Any]], ...] = (LessThan, LessThanOrEqual)
+    for operator in operators:
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: all nan column doesn't contain number"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: 1 is smaller than lower bound"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 10)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 10)).eval(data_file_nan)
         assert should_read, "Should match: 10 is larger than lower bound"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("min_max_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("min_max_nan", 1)).eval(data_file_nan)
         assert should_read, "Should match: no visibility"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan_null_bounds", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan_null_bounds", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: all nan column doesn't contain number"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: 1 is smaller than lower bound"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 10)).eval(  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 10)).eval(
             data_file_nan
         )
         assert should_read, "Should match: 10 larger than lower bound"
@@ -711,31 +714,32 @@ def test_inclusive_metrics_evaluator_less_than_and_less_than_equal(schema_data_f
 def test_inclusive_metrics_evaluator_greater_than_and_greater_than_equal(
     schema_data_file_nan: Schema, data_file_nan: DataFile
 ) -> None:
-    for operator in [GreaterThan, GreaterThanOrEqual]:
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+    operators: tuple[type[LiteralPredicate[Any]], ...] = (GreaterThan, GreaterThanOrEqual)
+    for operator in operators:
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: all nan column doesn't contain number"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 1)).eval(data_file_nan)
         assert should_read, "Should match: upper bound is larger than 1"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 10)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("max_nan", 10)).eval(data_file_nan)
         assert should_read, "Should match: upper bound is larger than 10"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("min_max_nan", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("min_max_nan", 1)).eval(data_file_nan)
         assert should_read, "Should match: no visibility"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan_null_bounds", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan_null_bounds", 1)).eval(data_file_nan)
         assert not should_read, "Should not match: all nan column doesn't contain number"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 1)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 1)).eval(data_file_nan)
         assert should_read, "Should match: 1 is smaller than upper bound"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 10)).eval(  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("some_nan_correct_bounds", 10)).eval(
             data_file_nan
         )
         assert should_read, "Should match: 10 is smaller than upper bound"
 
-        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 30)).eval(data_file_nan)  # type: ignore[arg-type]
+        should_read = _InclusiveMetricsEvaluator(schema_data_file_nan, operator("all_nan", 30)).eval(data_file_nan)
         assert not should_read, "Should not match: 30 is greater than upper bound"
 
 
@@ -1162,7 +1166,7 @@ def test_strict_missing_stats(strict_data_file_schema: Schema, strict_data_file_
         upper_bounds=None,
     )
 
-    expressions = [
+    expressions: list[BooleanExpression] = [
         LessThan("no_stats", 5),
         LessThanOrEqual("no_stats", 30),
         EqualTo("no_stats", 70),
@@ -1185,7 +1189,7 @@ def test_strict_zero_record_file_stats(strict_data_file_schema: Schema) -> None:
         file_path="file_1.parquet", file_format=FileFormat.PARQUET, partition=Record(), record_count=0
     )
 
-    expressions = [
+    expressions: list[BooleanExpression] = [
         LessThan("no_stats", 5),
         LessThanOrEqual("no_stats", 30),
         EqualTo("no_stats", 70),

tests/expressions/test_expressions.py

@@ -50,6 +50,7 @@
     IsNull,
     LessThan,
     LessThanOrEqual,
+    LiteralPredicate,
     Not,
     NotEqualTo,
     NotIn,
@@ -64,7 +65,7 @@
 from pyiceberg.expressions.literals import Literal, literal
 from pyiceberg.expressions.visitors import _from_byte_buffer
 from pyiceberg.schema import Accessor, Schema
-from pyiceberg.typedef import Record
+from pyiceberg.typedef import L, Record
 from pyiceberg.types import (
     DecimalType,
     DoubleType,
@@ -935,7 +936,7 @@ def test_bound_less_than_or_equal(term: BoundReference[Any]) -> None:
 
 def test_equal_to() -> None:
     equal_to = EqualTo(Reference("a"), literal("a"))
-    assert equal_to.model_dump_json() == '{"term":"a","type":"eq","value":"Any"}'
+    assert equal_to.model_dump_json() == '{"term":"a","type":"eq","value":"a"}'
     assert str(equal_to) == "EqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert repr(equal_to) == "EqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert equal_to == eval(repr(equal_to))
@@ -944,7 +945,7 @@ def test_equal_to() -> None:
 
 def test_not_equal_to() -> None:
     not_equal_to = NotEqualTo(Reference("a"), literal("a"))
-    assert not_equal_to.model_dump_json() == '{"term":"a","type":"not-eq","value":"Any"}'
+    assert not_equal_to.model_dump_json() == '{"term":"a","type":"not-eq","value":"a"}'
     assert str(not_equal_to) == "NotEqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert repr(not_equal_to) == "NotEqualTo(term=Reference(name='a'), literal=literal('a'))"
     assert not_equal_to == eval(repr(not_equal_to))
@@ -953,7 +954,7 @@ def test_not_equal_to() -> None:
 
 def test_greater_than_or_equal_to() -> None:
     greater_than_or_equal_to = GreaterThanOrEqual(Reference("a"), literal("a"))
-    assert greater_than_or_equal_to.model_dump_json() == '{"term":"a","type":"gt-eq","value":"Any"}'
+    assert greater_than_or_equal_to.model_dump_json() == '{"term":"a","type":"gt-eq","value":"a"}'
     assert str(greater_than_or_equal_to) == "GreaterThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert repr(greater_than_or_equal_to) == "GreaterThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert greater_than_or_equal_to == eval(repr(greater_than_or_equal_to))
@@ -962,7 +963,7 @@ def test_greater_than_or_equal_to() -> None:
 
 def test_greater_than() -> None:
     greater_than = GreaterThan(Reference("a"), literal("a"))
-    assert greater_than.model_dump_json() == '{"term":"a","type":"gt","value":"Any"}'
+    assert greater_than.model_dump_json() == '{"term":"a","type":"gt","value":"a"}'
     assert str(greater_than) == "GreaterThan(term=Reference(name='a'), literal=literal('a'))"
     assert repr(greater_than) == "GreaterThan(term=Reference(name='a'), literal=literal('a'))"
     assert greater_than == eval(repr(greater_than))
@@ -971,7 +972,7 @@ def test_greater_than() -> None:
 
 def test_less_than() -> None:
     less_than = LessThan(Reference("a"), literal("a"))
-    assert less_than.model_dump_json() == '{"term":"a","type":"lt","value":"Any"}'
+    assert less_than.model_dump_json() == '{"term":"a","type":"lt","value":"a"}'
     assert str(less_than) == "LessThan(term=Reference(name='a'), literal=literal('a'))"
     assert repr(less_than) == "LessThan(term=Reference(name='a'), literal=literal('a'))"
     assert less_than == eval(repr(less_than))
@@ -980,7 +981,7 @@ def test_less_than() -> None:
 
 def test_less_than_or_equal() -> None:
     less_than_or_equal = LessThanOrEqual(Reference("a"), literal("a"))
-    assert less_than_or_equal.model_dump_json() == '{"term":"a","type":"lt-eq","value":"Any"}'
+    assert less_than_or_equal.model_dump_json() == '{"term":"a","type":"lt-eq","value":"a"}'
     assert str(less_than_or_equal) == "LessThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert repr(less_than_or_equal) == "LessThanOrEqual(term=Reference(name='a'), literal=literal('a'))"
     assert less_than_or_equal == eval(repr(less_than_or_equal))
@@ -989,12 +990,12 @@ def test_less_than_or_equal() -> None:
 
 def test_starts_with() -> None:
     starts_with = StartsWith(Reference("a"), literal("a"))
-    assert starts_with.model_dump_json() == '{"term":"a","type":"starts-with","value":"Any"}'
+    assert starts_with.model_dump_json() == '{"term":"a","type":"starts-with","value":"a"}'
 
 
 def test_not_starts_with() -> None:
     not_starts_with = NotStartsWith(Reference("a"), literal("a"))
-    assert not_starts_with.model_dump_json() == '{"term":"a","type":"not-starts-with","value":"Any"}'
+    assert not_starts_with.model_dump_json() == '{"term":"a","type":"not-starts-with","value":"a"}'
 
 
 def test_bound_reference_eval(table_schema_simple: Schema) -> None:
@@ -1235,7 +1236,12 @@ def test_bind_ambiguous_name() -> None:
 #  |_|  |_|\_, |_|  \_, |
 #          |__/     |__/
 
-assert_type(EqualTo("a", "b"), EqualTo[str])
+
+def _assert_literal_predicate_type(expr: LiteralPredicate[L]) -> None:
+    assert_type(expr, LiteralPredicate[L])
+
+
+_assert_literal_predicate_type(EqualTo("a", "b"))
 assert_type(In("a", ("a", "b", "c")), In[str])
 assert_type(In("a", (1, 2, 3)), In[int])
 assert_type(NotIn("a", ("a", "b", "c")), NotIn[str])