Skip to content

Commit f6a3200

Browse files
committed
feat: make Or json serializable via IcebergBaseModel
1 parent 5ee5eea commit f6a3200

File tree

2 files changed

+31
-4
lines changed

2 files changed

+31
-4
lines changed

pyiceberg/expressions/__init__.py

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from __future__ import annotations
1919

20+
import typing
2021
from abc import ABC, abstractmethod
2122
from functools import cached_property
2223
from typing import (
@@ -32,14 +33,16 @@
3233
Union,
3334
)
3435

36+
from pydantic import ConfigDict, Field, field_serializer
37+
3538
from pyiceberg.expressions.literals import (
3639
AboveMax,
3740
BelowMin,
3841
Literal,
3942
literal,
4043
)
4144
from pyiceberg.schema import Accessor, Schema
42-
from pyiceberg.typedef import IcebergRootModel, L, StructProtocol
45+
from pyiceberg.typedef import IcebergBaseModel, IcebergRootModel, L, StructProtocol
4346
from pyiceberg.types import DoubleType, FloatType, NestedField
4447
from pyiceberg.utils.singleton import Singleton
4548

@@ -290,12 +293,18 @@ def __getnewargs__(self) -> Tuple[BooleanExpression, BooleanExpression]:
290293
return (self.left, self.right)
291294

292295

293-
class Or(BooleanExpression):
296+
class Or(IcebergBaseModel, BooleanExpression):
294297
"""OR operation expression - logical disjunction."""
295298

299+
model_config = ConfigDict(arbitrary_types_allowed=True)
300+
301+
type: str = Field(default="or", repr=False)
296302
left: BooleanExpression
297303
right: BooleanExpression
298304

305+
def __init__(self, left: typing.Union[BooleanExpression, Or], right: typing.Union[BooleanExpression, Or], *rest: Any):
306+
return super().__init__(left=left, right=right)
307+
299308
def __new__(cls, left: BooleanExpression, right: BooleanExpression, *rest: BooleanExpression) -> BooleanExpression: # type: ignore
300309
if rest:
301310
return _build_balanced_tree(Or, (left, right, *rest))
@@ -307,10 +316,24 @@ def __new__(cls, left: BooleanExpression, right: BooleanExpression, *rest: Boole
307316
return left
308317
else:
309318
obj = super().__new__(cls)
310-
obj.left = left
311-
obj.right = right
312319
return obj
313320

321+
@field_serializer("left")
322+
def ser_left(self, left: BooleanExpression) -> str:
323+
if isinstance(left, IcebergRootModel):
324+
return left.root
325+
return str(left)
326+
327+
@field_serializer("right")
328+
def ser_right(self, right: BooleanExpression) -> str:
329+
if isinstance(right, IcebergRootModel):
330+
return right.root
331+
return str(right)
332+
333+
def __str__(self) -> str:
334+
"""Return the string representation of the Or class."""
335+
return f"{str(self.__class__.__name__)}(left={repr(self.left)}, right={repr(self.right)})"
336+
314337
def __eq__(self, other: Any) -> bool:
315338
"""Return the equality of two instances of the Or class."""
316339
return self.left == other.left and self.right == other.right if isinstance(other, Or) else False

tests/expressions/test_expressions.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -718,6 +718,10 @@ def test_or() -> None:
718718
# Some syntactic sugar
719719
assert or_ == null | nan
720720

721+
assert (
722+
or_.model_dump_json()
723+
== '{"type":"or","left":"IsNull(term=Reference(name=\'a\'))","right":"IsNaN(term=Reference(name=\'b\'))"}'
724+
)
721725
assert str(or_) == f"Or(left={str(null)}, right={str(nan)})"
722726
assert repr(or_) == f"Or(left={repr(null)}, right={repr(nan)})"
723727
assert or_ == eval(repr(or_))

0 commit comments

Comments
 (0)