Awkward dunders for pow(x, y, z)

I have been recently working on a library python-flint that provides a type fmpz which is a faster version of Python’s int type similar to gmpy2’s `mpz type:

One thing that I needed to implement was 3-arg pow as in:

>>> from flint import fmpz
>>> pow(fmpz(2), fmpz(3), fmpz(5))

What I have found doing this is that the __pow__ dunder that is used for this is awkward to implement and that there is no way from Python (rather than Cython) to extend pow with a dunder on the second or third argument if there are mixed types.

This is how __pow__ should be implemented in Python if wanting to support the 3-arg pow:


_val = lambda x: x.val if isinstance(x, Integer) else x

class Integer:

    def __init__(self, val):
        self.val = val

    def __repr__(self):
        return f'Integer({self.val})'

    # Here mod must be an optional argument
    def __pow__(self, other, mod=None):
        print(f'Called: Integer.__pow__({self}, {other}, {mod})')
        return Integer(pow(_val(self), _val(other), _val(mod)))

    def __rpow__(self, other):
        print(f'Called: Integer.__rpow__({self}, {other})')
        return Integer(pow(_val(other), _val(self)))

For the 2-arg case that gives e.g.:

>>> Integer(2) ** Integer(3)
Called: Integer.__pow__(Integer(2), Integer(3), None)
>>> Integer(2) ** 3
Called: Integer.__pow__(Integer(2), 3, None)
>>> 2 ** Integer(3)
Called: Integer.__rpow__(Integer(3), 2)
>>> pow(2, Integer(3))
Called: Integer.__rpow__(Integer(3), 2)

Note that the third argument to __pow__ is not actually passed so it needs to be an optional argument like mod=None for this to work.

For the 3-arg case we have:

>>> pow(Integer(2), Integer(3), Integer(5))
Called: Integer.__pow__(Integer(2), Integer(3), Integer(5))
>>> pow(Integer(2), 3, 5)
Called: Integer.__pow__(Integer(2), 3, 5)
>>> pow(2, Integer(3), Integer(5))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for pow(): 'int', 'Integer', 'Integer'

What happens here is that in the 3-arg case pow(x, y, z) will always call type(x).__pow__(x, y, z) and there is no way for the other arguments y and z to influence this.

It is however possible to do this using Cython:

# integer.pyx

_val = lambda x: x.x if isinstance(x, CInteger) else x

cdef class CInteger:

    cdef object _x

    def __init__(self, x):
        self._x = x

    def x(self):
        return self._x

    def __repr__(self):
        return f'CInteger({self.x})'

    # Here mod must be a required argument
    def __pow__(self, other, mod):
        print(f'CInteger.__pow__({self}, {other}, {mod})')
        return CInteger(pow(_val(self), _val(other), _val(mod)))

    def __rpow__(self, other):
        print(f'CInteger.__rpow__({self}, {other})')
        return CInteger(pow(_val(other), _val(self)))

Then we can pass a CInteger as any of the 3 arguments to pow:

>>> pow(CInteger(2), CInteger(3), CInteger(5))
CInteger.__pow__(CInteger(2), CInteger(3), CInteger(5))
>>> pow(2, CInteger(3), CInteger(5))
CInteger.__pow__(2, CInteger(3), CInteger(5))
>>> pow(2, 3, CInteger(5))
CInteger.__pow__(2, 3, CInteger(5))
>>> pow(2, CInteger(3), 5)
CInteger.__pow__(2, CInteger(3), 5)

I got a bit lost in the CPython code trying to trace the whole codepath that make this happen but in the end pow(x, y, z) can end up calling type(y).__pow__(x, y, z) or type(z).__pow__(x, y, z) but there does not seem to be any way to implement the same in plain Python.

I think if I was designing this from scratch then I would not use the same __pow__ dunder for both x**y and pow(x, y, z) and instead have something like __pow3__. Of course there are three arguments to potentially dispatch on so it wouldn’t work just to have __pow3__ and __rpow3__. Perhaps instead it could work like in Cython where the same method is just called with differently typed arguments (so that self is not necessarily the expected type).

    def __pow3__(base, exp, mod):
        return Integer(pow(_val(base), _val(exp), _val(mod)))

Then pow(x, y, z) could in turn try:

  1. type(x).__pow3__(x, y, z)
  2. type(y).__pow3__(x, y, z)
  3. type(z).__pow3__(x, y, z)

I am not sure whether or not it is worth it to add a dunder methods like this. There can be more uses of pow(x, y, z) than just integers to be clear e.g. the same reasons that it makes sense to have a special operation for computing x**y mod z apply to polynomials and many other things besides integers. I am not aware of any libraries using 3-arg pow for anything apart from integers though.

The main thing that I dislike about the current situation is just that it does not seem to be possible in plain Python to make a type for y or z that can overload the behaviour of pow(int(x), y, z) when that is a feature that is being used by some popular python packages like gmpy2 already.

It looks like the code in CPython is made complicated by needing to support this ternary dispatch so maybe there is some better way to eliminate that altogether. I might be misunderstanding this but it looks like there is ternary dispatch code that exists purely for the case of handling 3-arg pow:

I imagine that it is also very difficult for other implementations like PyPy to match CPython’s behaviour here.

Perhaps there should instead be some way for other types to signify that they would like to coerce integers in mixed operations like:

class Integer:
    def __coerce__(self, other):
        if isinstance(other, int):
            return Integer(int_value)
            return NotImplemented

Then it would be possible for the pow builtin to coerce the types in some defined order and then call the appropriate __pow__ (or __pow3__) method. There are lots of other situations where being able to coerce types systematically would be useful.

Interestring write-up.

This would be so much easier if multiple dispatch were natively used for binary and ternary operators.

Agreed. The approach of having e.g. __add__ and __radd__ for “coperative dispatch” is awkward and error-prone compared to proper multiple dispatch. The case of 3-arg pow shows where this cooperative dispatch gets pushed clearly beyond its limits.

1 Like

I think this ship has sailed. We can definitely not change the rules around __add__ and __radd__ (and similar) without drastic impact on backward incompatibility, which we are not willing to do.

For 3-arg pow(), I believe the use case is not important enough for the complications that were introduced to __pow__ relative to the other binary operators, decades ago. Again, backwards compatibility makes it essentially impossible to change, but if we could change anything there, I would “solve” the problem by deprecating 3-arg pow().

1 Like

To be clear there is an open bug report in python-flint related to this with the most recent release of PyPy (7.3.12):

The fact that this feature (pow(2, fmpz(3), 5)) can only be implemented via the C API and (as far as I can tell) only by depending on subtle details of CPython’s type structs probably explains why it is hard to make something like this work across different Python implementations.

Potentially that bug report could be considered just a PyPy bug simply for PyPy not being compatible with CPython but there is no documentation anywhere that explains how any of this is supposed to work so it’s hard to say that PyPy is doing anything incorrect. Alternatively maybe it is a python-flint or Cython bug for depending on some subtle CPython details that were not supposed to be depended on. I am sure that other widely used code (at least gmpy2) depends on this as well though.