feat: make Or json serializable via IcebergBaseModel

Author: jaimeferj

pyiceberg/expressions/__init__.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import typing
 from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import (
@@ -35,6 +36,8 @@
 
 from pydantic import Field
 
+from pydantic import ConfigDict, Field, field_serializer
+
 from pyiceberg.expressions.literals import (
     AboveMax,
     BelowMin,
@@ -302,12 +305,18 @@ def __getnewargs__(self) -> Tuple[BooleanExpression, BooleanExpression]:
         return (self.left, self.right)
 
 
-class Or(BooleanExpression):
+class Or(IcebergBaseModel, BooleanExpression):
     """OR operation expression - logical disjunction."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    type: str = Field(default="or", repr=False)
     left: BooleanExpression
     right: BooleanExpression
 
+    def __init__(self, left: typing.Union[BooleanExpression, Or], right: typing.Union[BooleanExpression, Or], *rest: Any):
+        return super().__init__(left=left, right=right)
+
     def __new__(cls, left: BooleanExpression, right: BooleanExpression, *rest: BooleanExpression) -> BooleanExpression:  # type: ignore
         if rest:
             return _build_balanced_tree(Or, (left, right, *rest))
@@ -319,10 +328,24 @@ def __new__(cls, left: BooleanExpression, right: BooleanExpression, *rest: Boole
             return left
         else:
             obj = super().__new__(cls)
-            obj.left = left
-            obj.right = right
             return obj
 
+    @field_serializer("left")
+    def ser_left(self, left: BooleanExpression) -> str:
+        if isinstance(left, IcebergRootModel):
+            return left.root
+        return str(left)
+
+    @field_serializer("right")
+    def ser_right(self, right: BooleanExpression) -> str:
+        if isinstance(right, IcebergRootModel):
+            return right.root
+        return str(right)
+
+    def __str__(self) -> str:
+        """Return the string representation of the Or class."""
+        return f"{str(self.__class__.__name__)}(left={repr(self.left)}, right={repr(self.right)})"
+
     def __eq__(self, other: Any) -> bool:
         """Return the equality of two instances of the Or class."""
         return self.left == other.left and self.right == other.right if isinstance(other, Or) else False

tests/expressions/test_expressions.py

@@ -702,6 +702,10 @@ def test_or() -> None:
     # Some syntactic sugar
     assert or_ == null | nan
 
+    assert (
+        or_.model_dump_json()
+        == '{"type":"or","left":"IsNull(term=Reference(name=\'a\'))","right":"IsNaN(term=Reference(name=\'b\'))"}'
+    )
     assert str(or_) == f"Or(left={str(null)}, right={str(nan)})"
     assert repr(or_) == f"Or(left={repr(null)}, right={repr(nan)})"
     assert or_ == eval(repr(or_))