CODEWITHSUNDEEP
pythonpython-3.x
Blaine
Blaine

Reputation: 677

Magic methods in python and the __str__ function

I am trying to implement my own class for complex numbers, to better understand how classes work in python. I have been trying to replace the str magic method to print the complex number in the a+bi format.

def __str__(self):
    out="%i" % self.real
    if self.imaginary==0: return out
    if self.imaginary>=0: out+="+%ii" % self.imaginary
    else: out+="%ii" % self.imaginary
    return out

what I am interested in is the pythonic way of writing this unappealing block of code, if there is any to implement the fact that if imaginary part is negative, i should get a-bi and if imaginary part is 0 i should get a?

Upvotes: 1

Views: 2767

Answers (2)

norok2
norok2

Reputation: 26946

What about just going a bit more explicit:

class Complex(object):
    IM_TOKEN = 'i'

    def __init__(self, real_part=0, imag_part=0):
        self.real = real_part
        self.imag = imag_part

    def null_imag(self):
        return self.imag == 0

    def null_real(self):
        return self.real == 0

    def __str__(self):
        if self.null_imag():
            return str(self.real)
        elif self.null_real():
            return str(self.imag) + type(self).IM_TOKEN
        else:
            return '{}{:+}{}'.format(self.real, self.imag, type(self).IM_TOKEN)

and testing it:

import itertools

for r, i in itertools.product([1, -1, 0], [1, -1, 0]):
    print(f'real: {r}, imag: {i}, complex: {Complex(r, i)}')

real: 1, imag: 1, complex: 1+1i
real: 1, imag: -1, complex: 1-1i
real: 1, imag: 0, complex: 1
real: -1, imag: 1, complex: -1+1i
real: -1, imag: -1, complex: -1-1i
real: -1, imag: 0, complex: -1
real: 0, imag: 1, complex: 1i
real: 0, imag: -1, complex: -1i
real: 0, imag: 0, complex: 0

As a side note the == does not work well with float data.

Upvotes: 1

lmiguelvargasf
lmiguelvargasf

Reputation: 69983

If you are using Python 3.6+, use f-strings as follows:

def __str__(self):
    if self.imaginary == 0:
        return f'{self.real}'
    if self.real == 0:
        return f'{self.imaginary}i'

    return f'{self.real} {self.imaginary:+}i'

If you are using a version prior to Python 3.6, you should use format.

def __str__(self):
    if self.imaginary == 0: return '{}'.format(self.real)
    if self.real == 0: return '{}i'.format(self.imaginary)

    return '{} {:+}i'.format(self.real, self.imaginary)

I have also improved a bit the logic. Basically, when it has no imaginary part it is just returning the real part, if there is no real part, it is returning the imaginary part. When it has both imaginary and real part, it returns the complex number.

Notice that + specified after the two dots is the format. This format allows you to have the the sign of the imaginary part printed.

Examples

Assuming that your class is named CNumber

>>> x = CNumber(10, 1)
>>> str(x)
'10 +1i'
>>> x = CNumber(5, 0)
>>> str(x)
'5'
>>> x = CNumber(0, 3)
'3i'
>>> x = CNumber(1, -1)
'1 -1i'

If you want the following format a + bi

def __str__(self):
    if self.imaginary == 0: return f'{self.real}'
    if self.real == 0: return f'{self.imaginary}i'

    sign = '+' if self.imaginary > 0 else '-'

    return f'{self.real} {sign} {abs(self.imaginary)}i'

Upvotes: 1

Related Questions