Skip to content

Commit

Permalink
feat(api): implement all bitwise operators
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and kszucs committed Sep 5, 2022
1 parent ff1a680 commit 7fc5073
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 0 deletions.
38 changes: 38 additions & 0 deletions ibis/expr/operations/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -316,3 +316,41 @@ class Sin(TrigonometricUnary):
@public
class Tan(TrigonometricUnary):
"""Returns the tangent of x"""


@public
class BitwiseNot(Unary):
arg = rlz.integer
output_dtype = rlz.numeric_like("args", operator.invert)


class BitwiseBinary(Binary):
left = rlz.integer
right = rlz.integer


@public
class BitwiseAnd(BitwiseBinary):
output_dtype = rlz.numeric_like("args", operator.and_)


@public
class BitwiseOr(BitwiseBinary):
output_dtype = rlz.numeric_like("args", operator.or_)


@public
class BitwiseXor(BitwiseBinary):
output_dtype = rlz.numeric_like("args", operator.xor)


@public
class BitwiseLeftShift(BitwiseBinary):
output_shape = rlz.shape_like("args")
output_dtype = dt.int64


@public
class BitwiseRightShift(BitwiseBinary):
output_shape = rlz.shape_like("args")
output_dtype = dt.int64
70 changes: 70 additions & 0 deletions ibis/expr/types/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from public import public

from ibis.common.exceptions import IbisTypeError
from ibis.expr.types.core import _binop
from ibis.expr.types.generic import Column, Scalar, Value

Expand Down Expand Up @@ -765,6 +766,75 @@ def convert_base(

return ops.BaseConvert(self, from_base, to_base).to_expr()

def __and__(self, other: IntegerValue) -> IntegerValue | NotImplemented:
"""Bitwise and `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseAnd, self, other)

__rand__ = __and__

def __or__(self, other: IntegerValue) -> IntegerValue | NotImplemented:
"""Bitwise or `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseOr, self, other)

__ror__ = __or__

def __xor__(self, other: IntegerValue) -> IntegerValue | NotImplemented:
"""Bitwise xor `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseXor, self, other)

__rxor__ = __xor__

def __lshift__(self, other: IntegerValue) -> IntegerValue | NotImplemented:
"""Bitwise left shift `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseLeftShift, self, other)

def __rlshift__(
self, other: IntegerValue
) -> IntegerValue | NotImplemented:
"""Bitwise left shift `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseLeftShift, other, self)

def __rshift__(self, other: IntegerValue) -> IntegerValue | NotImplemented:
"""Bitwise right shift `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseRightShift, self, other)

def __rrshift__(
self, other: IntegerValue
) -> IntegerValue | NotImplemented:
"""Bitwise right shift `self` with `other`."""
from ibis.expr import operations as ops

return _binop(ops.BitwiseRightShift, other, self)

def __invert__(self) -> IntegerValue:
"""Bitwise not of `self`.
Returns
-------
IntegerValue
Inverted bits of `self`.
"""
from ibis.expr import operations as ops

try:
node = ops.BitwiseNot(self)
except (IbisTypeError, NotImplementedError):
return NotImplemented
else:
return node.to_expr()


@public
class IntegerScalar(NumericScalar, IntegerValue):
Expand Down

0 comments on commit 7fc5073

Please sign in to comment.