z3py compare Datatype / Enum to string

Following this example in (found here: z3py) I can compare c to e.g. Color.green.

Color = Datatype('Color')
Color.declare('red')
Color.declare('green')
Color.declare('blue')
Color = Color.create()

# Let c be a constant of sort Color
c = Const('c', Color)
# Then, c must be red, green or blue
prove(Or(c == Color.green, 
         c == Color.blue,
         c == Color.red))

In my application I have to compare c to a python-string: I would like something like this:

c = Const('c', Color)
solve(c == "green") # this doesn't work, but it works with Color.green

The approach works e.g. for IntSort (see below), but not for my own Datatype.

i = Int("i")
solve(i < 10)

Solution

One solution that worked for me (comparing Datatypes/Enums to strings) was to add a cast routine to class DatatypeSortRef(SortRef) in z3.py. It will try to find a constructor that matches the given string and use it, otherwise continue with existing behaviour (super().cast(val))

Here's the code that I used:

def cast(self, val):
    """This is so we can cast a string to a Z3 DatatypeRef. This is useful if we want to compare strings with a Datatype/Enum to a String.

    >>> Color = Datatype("Color")
    >>> Color.declare("red")
    >>> Color.declare("green")
    >>> Color.declare("blue")
    >>> Color = Color.create()

    >>> x = Const("x", Color)
    >>> solve(x != "red", x != "blue")
    [x = green]
    """
    if type(val) == str:
        for i in range(self.num_constructors()):
            if self.constructor(i).name() == val:
                return self.constructor(i)()
    return super().cast(val)

Note: I did not pay attention to general correctness. This approach works for me, but might cause problems with your code.