overload builtin operator

[Bengt Richter]

On Thu, 25 Aug 2005 16:12:20 +0200, Reinhold Birkenfeld <reinhold-birkenfeld-nospam at wolke7.net> wrote:

>Shaun wrote:
>> Hi,
>> 
>> I'm trying to overload the divide operator in python for basic arithmetic.
>> eg. 10/2 ... no classes involved.
>> 
>> I am attempting to redefine operator.__div__ as follows:
>> 
>>      # my divide function
>>      def safediv(a,b):
>>          return ...
>> 
>>      # reassign buildin __div__
>>      import operator
>>      operator.__div__ = safediv
>> 
>> The operator.__dict__ seems to be updated OK but the '/' operator still  
>> calls buildin __div__
>
>It won't work that way. You cannot globally modify the behaviour of an operator,
>but you can customize how an operator works for your type.
>
>Consider:
>
>class safeint(int):
>    def __div__(self, other):
>        return safediv(self, other)
>
>safeint(10)/2
>
You are right that you cannot globally modify the behaviour of an operator in the
sense the OP seems to be envisioning, but with some trouble I think it would be possible
to interfere with the translation of '<numerator-term>/<denominator-term>' to become
'safediv(<numerator-term>, <denominator-term>)' by transforming the AST during a custom
import process, such that wherever a Div node is found, a CallFunc node is substituted. E.g.,

for a node like

    Div((Name('numerator'), Name('denominator')))

substitute another node like

    CallFunc(Name('safediv'), [Name('numerator'), Name('denominator')], None, None)

where the Name('numerator') and Name('denominator') in the latter are actually
the Div node's .left and .right, so as to operate on the same args (which can
be abitrary expression-representing subtrees).

Of course, you probably also want to replace

    AugAssign(Name('a'), '/=', Name('b'))

with

    Assign([AssName('a', 'OP_ASSIGN')], CallFunc(Name('safediv'), [Name('a'), Name('b')], None, None))

Unfortunately, the AST tree does not seem to be designed to be edited easily wrt
_replacing_ nodes found by walking via flattened lists vs. just _using_ them in order.

I think I can create special walker that can do replacements of nodes found anywhere, and
calling node-type-specific or default supplied callbacks to supply replacements for original nodes
passed to them, but this will require a number of specialized visitor methods etc., so I will
have to get back to this later. But at least the OP and you nudged me into thinking about
AST rewriting again ;-) (also affects my @@sourcedeco ideas ;-)

... this didn't go out due to news server problem, so I'll add some clips from what I did re ast editing:

----< testdiv.py >-------------------------------------------------------------------------
def test():
    print 1.0/2.0
    print 12/3
    a=12; b=3
    print a/b
    print 2**a/(b+1)
    try:
        print 'print 1/0 =>'
        print 1/0
    except Exception, e: print '%s: %s' %(e.__class__.__name__, e)
    try:
        print 'print a/(b*(a-12)) =>'
        print a/(b*(a-12))
    except Exception, e: print '%s: %s' %(e.__class__.__name__, e)
    try:
        print 'def foo(x=1/(b-3)): return x =>'
        def foo(x=1/(b-3)): return x
        print 'print foo() =>'
        print foo()
    except Exception, e: print '%s: %s' %(e.__class__.__name__, e)
    try:
        print 'b /= (a-12) =>'
        b /= (a-12)
        print 'b after b/=(a-12):', b
    except Exception, e: print '%s: %s' %(e.__class__.__name__, e)
    
if __name__ == '__main__': test()
-------------------------------------------------------------------------------------------
Outputfrom run without conversion of / :

0.5
4
4
1024
print 1/0 =>
ZeroDivisionError: integer division or modulo by zero
print a/(b*(a-12)) =>
ZeroDivisionError: integer division or modulo by zero
def foo(x=1/(b-3)): return x =>
ZeroDivisionError: integer division or modulo by zero
b /= (a-12) =>
ZeroDivisionError: integer division or modulo by zero


----< import_safediv.py >------------------------------------------------------------------
# import_safediv.py
from importast import importast
from compiler.ast import Div, CallFunc, Name
from compiler.ast import AugAssign, Assign, AssName

def safediv(num, den):
    if den==0: result = 0*num
    else: result = num/den
    print 'safediv(%r, %r) => %r'%(num, den, result)
    return result

def div2safediv(divnode):
    """replace Div nodes with CallFunc nodes calling safediv with same args"""
    return  CallFunc(Name('safediv'),[divnode.left, divnode.right], None, None, divnode.lineno)

def divaugass2safediv(auganode):
    if auganode.op != '/=': return auganode
    assname =auganode.node.name
    return Assign([AssName(assname, 'OP_ASSIGN')],
        CallFunc(Name('safediv'),[Name(assname), auganode.expr], None, None, auganode.lineno))

callbacks = {Div:div2safediv, AugAssign:divaugass2safediv}
def import_safediv(modname, verbose=False):
    return importast(modname, callbacks, verbose)

if __name__ == '__main__':
    import sys
    modname, args = sys.argv[1], sys.argv[2:]
    verbose = (args[0:] and args[0]=='-verbose') or False
    modsafe = import_safediv(modname, verbose)
    modsafe.safediv = safediv
    if hasattr(modsafe, 'test'): modsafe.test()
-------------------------------------------------------------------------------------------

Result from running the above and specifying testdiv as the module
to import and enforce safediv on, and run test() on:

safediv(1.0, 2.0) => 0.5
0.5
safediv(12, 3) => 4
4
safediv(12, 3) => 4
4
safediv(4096, 4) => 1024
1024
print 1/0 =>
safediv(1, 0) => 0
0
print a/(b*(a-12)) =>
safediv(12, 0) => 0
0
def foo(x=1/(b-3)): return x =>
safediv(1, 0) => 0
print foo() =>
0
b /= (a-12) =>
safediv(3, 0) => 0
b after b/=(a-12): 0

Note that the exceptions went away ;-)

It was much messier than it should have been, and the code is not very efficient,
but at least it's a kind of proof of concept ;-)

Regards,
Bengt Richter