Overloading __eq__ to return custom objects

Question

I am writing a DSL in Python and I want to overload operators to be able to easily write expressions of my DSL. For example, I want to write Var("a") + Var("b") and get the equivalent representation of Add(Var("a"), Var("b")). For this, I overloaded the __add__ method, and it works fine for this one.

Nevertheless, I try to overload the __eq__ method to achieve something similar: I want to write Var("a") == Var("b") and get the equivalent representation of Eq(Var("a"), Var("b")). By overloading the __eq__ method, returning an instance of Eq, I achieved my goal. But when overloading the __eq__ method, it obviously interferes with the behaviour of standard Python, such as Var("b") in [Var("a")] returning True.

Is there a way to achieve my goal, that is, being able to write Var("a") == Var("b") and get Eq(Var("a"), Var("b")), but still being able to write if Var("a") == Var("b"): blablabla or putting expressions in built-in containers, etc.?

EDIT

I tried to implement the __bool__ method of Eq class, and it seems to work (see the following code). Is there something I'm missing or is it a viable solution?

class Expr:
    def __add__(self, other):
        return Add(self, other)

    def __eq__(self, other):
        return Eq(self, other)

    def __repr__(self):
        return str(self)

    def __add__(self, other):
        return Add(self, other)

    def __ne__(self, other):
        return Neq(self, other)

class Var(Expr):
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return "Var(" + str(self.name) + ")"

    def equals(self, other):
        if type(self) is type(other):
            return self.name == other.name
        else:
            return False

    def __hash__(self):
        return 17 + 23 * hash(self.name)

class Add(Expr):
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def __str__(self):
        return "Add(" + str(self.left) + ", " + str(self.right) + ")"

    def equals(self, other):
        if type(self) is type(other):
            return ( ( self.left.equals(other.left) and
                       self.right.equals(other.right) ) or
                     ( self.left.equals(other.right) and
                       self.right.equals(other.left) ) )
        else:
            return False

    def __hash__(self):
        return (17 + 23 * hash("+") +
                23 * 23 * hash(self.left) + 23 * 23 * hash(self.right))

class Eq(Expr):
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def __str__(self):
        return "Eq(" + str(self.left) + ", " + str(self.right) + ")"

    def equals(self, other):
        if type(self) is type(other):
            return ( ( self.left.equals(other.left) and
                       self.right.equals(other.right) ) or
                     ( self.left.equals(other.right) and
                       self.right.equals(other.left) ) )
        else:
            return False

    def __bool__(self):
        return self.left.equals(self.right)

    def __hash__(self):
        return (17 + 23 * hash("==") +
                23 * 23 * hash(self.left) + 23 * 23 * hash(self.right))

class Neq(Expr):
    def __init__(self, left, right):
        self.left = left
        self.right = right

    def __str__(self):
        return "Neq(" + str(self.left) + ", " + str(self.right) + ")"

    def equals(self, other):
        if type(self) is type(other):
            return ( ( not self.left.equals(other.left) or
                       not self.right.equals(other.right) ) and
                     ( not self.left.equals(other.right) or
                       not self.right.equals(other.left) ) )
        else:
            return False

    def __bool__(self):
        return not self.left.equals(self.right)

    def __hash__(self):
        return (17 + 23 * hash("!=") +
                23 * 23 * hash(self.left) + 23 * 23 * hash(self.right))


a = Var("a")
aa = Var("a")
b = Var("b")
c = Var("c")


print("a + b", "=>", a + b)   # a + b => Add(Var(a), Var(b))
print("a == b", "=>", a == b) # a == b => Eq(Var(a), Var(b))
print("a != b", "=>", a != b) # a != b => Neq(Var(a), Var(b))

print("a if a == b else b", "=>", a if a == b else b)
# a if a == b else b => Var(b)
print("a if a == aa else b", "=>", a if a == aa else b)
# a if a == aa else b => Var(a)


l = [a, a+b]
print("l", "=>", l)               # l => [Var(a), Add(Var(a), Var(b))]
print("b in l", "=>", b in l)     # b in l => False
print("a in l", "=>", a in l)     # a in l => True
print("aa in l", "=>", aa in l)   # aa in l => True
print("a+b in l", "=>", a+b in l) # a+b in l => True
print("b+a in l", "=>", b+a in l) # b+a in l => True
print("a+c in l", "=>", a+c in l) # a+c in l => False


if a == b:
    print("a == b is True")
else:
    print("a == b is False")        # a == b is False
if a == aa:
    print("a == aa is True")        # a == aa is True
else:
    print("a == aa is False")

if a != b:
    print("a != b is True")         # a != b is True
else:
    print("a != b is False")
if a != aa:
    print("a != aa is True")
else:
    print("a != aa is False")       # a != aa is False


if a == b or a == aa:
    print("a == b or a == aa is True")   # a == b or a == aa is True
else:
    print("a == b or a == aa is False")
if a == aa and a == b:
    print("a == aa and a == b is True")
else:
    print("a == aa and a == b is False") # a == aa and a == b is False
if not a == aa:
    print("not a == aa is True")
else:
    print("not a == aa is False")        # not a == aa is False
if not a == b:
    print("not a == b is True")          # not a == b is True
else:
    print("not a == b is False")


if a == 3:
    print("a == 3 is True")
else:
    print("a == 3 is False")             # a == 3 is False
if a != 3:
    print("a != 3 is True")              # a != 3 is True
else:
    print("a != 3 is False")
if 3 == a:
    print("3 == a is True")
else:
    print("3 == a is False")             # 3 == a is False
if 3 != a:
    print("3 != a is True")              # 3 != a is True
else:
    print("3 != a is False")


if a == 'a':
    print("a == 'a' is True")
else:
    print("a == 'a' is False")           # a == 'a' is False
if a != 'a':
    print("a != 'a' is True")            # a != 'a' is True
else:
    print("a != 'a' is False")
if 'a' == a:
    print("'a' == a is True")
else:
    print("'a' == a is False")           # 'a' == a is False
if 'a' != a:
    print("'a' != a is True")            # 'a' != a is True
else:
    print("'a' != a is False")


s = {a}
print("s", "=>", s)             # s => {Var(a)}
print("a in s", "=>", a in s)   # a in s => True
print("b in s", "=>", b in s)   # b in s => False
print("aa in s", "=>", aa in s) # aa in s => True

d = {a: 1, b: 2}
print("d", "=>", d)             # d => {Var(b): 2, Var(a): 1}
print("d[a]", "=>", d[a])       # d[a] => 1
print("d[b]", "=>", d[b])       # d[b] => 2
print("c in d", "=>", c in d)   # c in d => False
print("aa in d", "=>", aa in d) # aa in d => True
print("d[aa]", "=>", d[aa])     # d[aa] => 1

Overloading eq to return custom objects

Answers (1)

Related Questions