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))
3
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
:
# integer.py
_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(8)
>>> Integer(2) ** 3
Called: Integer.__pow__(Integer(2), 3, None)
Integer(8)
>>> 2 ** Integer(3)
Called: Integer.__rpow__(Integer(3), 2)
Integer(8)
>>> pow(2, Integer(3))
Called: Integer.__rpow__(Integer(3), 2)
Integer(8)
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))
Integer(3)
>>> pow(Integer(2), 3, 5)
Called: Integer.__pow__(Integer(2), 3, 5)
Integer(3)
>>> 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
@property
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))
CInteger(3)
>>> pow(2, CInteger(3), CInteger(5))
CInteger.__pow__(2, CInteger(3), CInteger(5))
CInteger(3)
>>> pow(2, 3, CInteger(5))
CInteger.__pow__(2, 3, CInteger(5))
CInteger(3)
>>> pow(2, CInteger(3), 5)
CInteger.__pow__(2, CInteger(3), 5)
CInteger(3)
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:
type(x).__pow3__(x, y, z)
type(y).__pow3__(x, y, z)
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)
else:
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.