Description
Hi Z3 folks,
I've been playing around with an idea for z3.BoolRef
and wanted to see what you all think. Basically, it's about adding some operator overloads to make things a bit more slick and readable.
Here's the gist of it:
from __future__ import annotations
import z3
class Bool(z3.BoolRef):
def __init__(self, arg: str | z3.BoolRef):
if isinstance(arg, str):
super().__init__(z3.Bool(arg).as_ast())
elif isinstance(arg, z3.BoolRef):
super().__init__(arg.as_ast())
else:
raise TypeError(f'Expected str or z3.BoolRef, got {type(arg)}')
def __and__(self, other: Bool) -> Bool:
return Bool(z3.And(self, other))
def __or__(self, other: Bool) -> Bool:
return Bool(z3.Or(self, other))
def __invert__(self) -> Bool:
return Bool(z3.Not(self))
def __rshift__(self, other: Bool) -> Bool:
return Bool(z3.Implies(self, other))
def __lshift__(self, other: Bool) -> Bool:
return Bool(z3.Implies(other, self))
def __add__(self, other: Bool) -> z3.ArithRef:
return z3.If(self, 1, 0) + z3.If(other, 1, 0)
This lets you write logical expressions more like how you'd do it in plain Python, e.g. ~(x & y) == (~x | ~y)
. Feels cleaner to me, but I'm curious about what others think.
If this looks like something useful, I'm thinking of putting together a PR. I'm guessing the changes would go somewhere around here in the code:
Lines 1568 to 1586 in 9382b96
But hey, I'm not entirely sure if that's the right spot or if there might be any unintended side effects, so any pointers or insights would be really helpful.