Search code examples
pythoncomplex-numbersiterable-unpacking

Can a complex number be unpacked into x and y in python?


I want to be able to unpack a complex number into its real and imaginary parts like so:

>>> z = 3 + 5j
>>> x, y = z
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: cannot unpack non-iterable complex object

but I get an error. Of course I can do:

>>> x, y = z.real, z.imag

but this seems repetitive and not very readable, I can improve this by writing a function like this:

def unpack(z: complex) -> tuple[float, float]:
    return z.real, z.imag
x, y = unpack(z)

but I feel like there should be a better way.

Do you know a better way to unpack a complex number into its real and imaginary parts?


Solution

  • One thing you can do is subclass complex and implement a custom __iter__ method that returns the unpacked version you seek.

    I chose to call it complez in this case, since z is often a symbol for a complex number.

    import cmath
    
    class complez(complex):
        def __iter__(self):
            return iter((self.real, self.imag))
    
    z = complez(3, 5)
    x, y = z
    print(x, y)
    # 3.0 5.0
    

    You can also do the following, if you don't mind overriding the existing complex. Using this approach allows you to continue using the complex class throughout your code without needing to make any changes to naming.

    import builtins
    import cmath
    
    class complex(builtins.complex):
        def __iter__(self):
            return iter((self.real, self.imag))
    
    z = complex(3, 5)
    x, y = z
    print(x, y)
    # 3.0 5.0
    

    (Technically, you could also drop the use of the builtins module).

    In either case, the standard unpacking operator would also work, e.g.

    print(*z)
    # 3.0 5.0
    

    NOTE: This solves your immediate problem, but it does have consequences. When you perform operations on these new numbers, e.g. adding 2 complez numbers, you would get a complex number result -- which won't have the convenient __iter__ defined, so you'd have to override other methods to account for this side-effect. For example:

    def __add__(self, other):
        return complez(super().__add__(other))
    

    Hence, this approach does result in more overhead in producing a fully general solution. But, this is the tradeoff for having the convenience of natural unpacking.