In Python, is there a way to automatically cast derived class operation results from an operator inherited from parent class to that derived class?

Viewed 71

Suppose I made a custom class from a parent class that has all kinds of operators (+,-,*,/,etc) defined, and I want to keep how those operators work as-is. However, without overriding the operators, the returned values from them are of the parent class type. Is there a convenient way to automatically cast the operator results to my derived class type?

Here's an example -- MySeries is a class based on Pandas' Series class.

import pandas as pd

class MySeries(pd.Series):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

if I make some instances from it, they'll be type MySeries, which is good.

s1 = MySeries([1,2,3])
s2 = MySeries([4,5,6])

type(s1)
>>> __main__.MySeries

But when I do an operation like +, the returned value becomes the parent class:

type(s1+s2)
>>> pandas.core.series.Series

My hope is type(s1+s2) can be of type MySeries. I know a dumb way to do it is to override e.g. the __add__() method:

class MySeries(pd.Series):
    def __add__(self, other):
        return self.__class__(super().__add__(self, other))

but doing the same for all operators is quite tedious.

Is there a more elegant way to do this?

1 Answers

This seems to do the trick:

class MySeries(pd.Series):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @property
    def _constructor(self):
        return MySeries

To find this, I stepped through the addition arithmetic call with a debugger in order to find where the new value was created. It eventually called pd.Series._constructor for this.

Related