Search code examples
pythonreturninstantiation

Return different type upon class instantiation?


So I am making my own version of a Complex class which works as a representation of imaginary numbers: Complex(a, b) where a + bi.

Thing is, I want to return something else based on arguments given. Im pretty sure I have to use the magic method new but I cant seem to be able to handle the arguments given. For instance:

a = Complex(4, 5)
print(a)
> Complex(4, 5)
b = Complex(3, 0)
print(b)
> 3
c = Complex(0, 0)
print(c)
> 0

I want the b and c variables to be assigned numerical values. Not some sort of to-string solution where the class pretends to be something it is not. How can I do this?

Thanks for replies :))


Solution

  • As others have pointed out in comments, this is exactly what __new__ is for.

    Here's what happens when you create an instance of a class Complex:

    1. You write Complex(a, b).
    2. Complex.__new__(a, b) is called, and by default (that is, if you don't override the __new__ method) this returns an empty object of type Complex.
    3. If the object returned by Complex.__new__(a, b) (let's denote it as cmplx) is actually of type Complex, call cmplx.__init__(a, b), or equivalently, Complex.__init__(cmplx, a, b).

    For your case, simply return an int in the __new__ method when the second argument is called, and call the __new__ method of the super class otherwise. That is:

    class Complex:
        def __new__(cls, a, b):
            if b == 0:
                return a
            return super().__new__(cls)  # or equiv.: object.__new__(cls)