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?
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.