Search code examples
pythonoopnumpymathmagic-methods

How do I make existing functions in math and numpy module support user-defined objects?


I defined a class that I'd like basic operations such as +, -, and even more advanced math operations such as sin to support. That is, I need to define some new rules for these basic operators and math functions.

I could deal with basic operators using Python's magic methods as below,

class NewObj():

    def __init__(self, *args):
        pass

    def __add__(self, other):
        pass

    def __mul__(self, other):
        pass

    # Use magic methods to define other basic operators below...

And I could also redefine basic functions in the same .py file where class NewObj is defined, such as

def sin(x):
    pass

And name this file as, say, myoperator.py. Then I could import this module and apply new operations on NewObj objects.


But I also want existing functions in numpy and math to support my NewObj objects so that math.sin() and numpy.sin() also support my newly defined NewObj objects. How can I achieve this?

Another question is: is it possible to encapsulate functions inside my NewObj class just like magic methods so that all the stuff are written inside a class data structure?


Solution

  • The math module explicitly documents that functions like math.sin always return floats. If you want math.sin(your_object) to return a NewObj instance instead of a float, don't do that. It'll confuse the hell out of everyone if you even make it work, and it'll cause initialization order bugs and other headaches. (There's a reason NumPy has its own numpy.sin instead of trying to make math.sin support NumPy arrays.)

    If you're okay with math.sin(your_object) returning a float, then implement a __float__ method to convert your objects to floats:

    class NewObj(object):
        ...
        def __float__(self):
            return whatever
    

    math.sin will convert your object to a float, compute the sine of the float, and return the sine as a float.


    For NumPy, just implement a sin method:

    class NewObj(object):
        ...
        def sin(self):
            return whatever
    

    numpy.sin will delegate to your sin method. You can have it return a NewObj; there's no need to cast to float or anything. Most similar NumPy functions will delegate to methods similarly. That said, trying to use custom objects inside NumPy arrays defeats most of the efficiency benefits of NumPy, so you might want to rethink your design if you want to do that.