Trimming zero elements at the end in a decorator

user_11751294 0
Viewed 41

To pretty print my calculation results I wrote a decorator:

from functools import wraps

def print_disp(f):
    @wraps(f)
    def wrapping(*args, **kwargs):
        disp, disp_std = f(*args, **kwargs)
        labels = ('GD', 'GDD','TOD', 'FOD', 'QOD')
        for i, (label, disp_item, disp_std_item) in enumerate(zip(labels, disp, disp_std)):
            print(f'{label} = {disp_item:.5f} ± {disp_std_item:.5f} fs^{i+1}')
        return disp, disp_std
    return wrapping

which can be used the following way:

import numpy as np

@print_disp
def calc():
    disp = np.array([1, 0, 7, 0, 0])
    disp_std = np.array([0.1, 0, 0.7, 0, 0])
    return disp, disp_std

and it returns

GD = 1.00000 ± 0.10000 fs^1
GDD = 0.00000 ± 0.00000 fs^2
TOD = 7.00000 ± 0.70000 fs^3
FOD = 0.00000 ± 0.00000 fs^4
QOD = 0.00000 ± 0.00000 fs^5

Which is fine. However, I want to avoid printing redundant 0 lines, but I only want to trim the zeros at the end, so the line GDD = 0 should stay. Because of this reason I couldn't add something like if disp item != 0: print(..) to the decorator body. The desired behaviour is:

GD = 1.00000 ± 0.10000 fs^1
GDD = 0.00000 ± 0.00000 fs^2
TOD = 7.00000 ± 0.70000 fs^3

How can I solve this?

python-decorators python
2 Answers

If you want to print only the first line with zeroes, you can modify the decorator like this to keep track of whether or not such a line is printed:

def print_disp(f):
    @wraps(f)
    def wrapping(*args, **kwargs):
        zero_printed = False
        disp, disp_std = f(*args, **kwargs)
        labels = ('GD', 'GDD','TOD', 'FOD', 'QOD')
        for i, (label, disp_item, disp_std_item) in enumerate(zip(labels, disp, disp_std)):
            if disp_item == 0 and zero_printed:
                continue
            zero_printed |= disp_item == 0
            print(f'{label} = {disp_item:.5f} ± {disp_std_item:.5f} fs^{i+1}')
        return disp, disp_std
    return wrapping
user_15772600

You can do something like this

from functools import wraps

def print_disp(f):
    @wraps(f)
    def wrapping(*args, **kwargs):
        disp, disp_std = f(*args, **kwargs)
        labels = ('GD', 'GDD','TOD', 'FOD', 'QOD')
        for i, (label, disp_item, disp_std_item) in enumerate(zip(labels, disp, disp_std)):
            if sum(disp[i:])+sum(disp_std[i:]) == 0:
              break
            print(f'{label} = {disp_item:.5f} ± {disp_std_item:.5f} fs^{i+1}')
            last_item = disp_item+disp_std_item
        return disp, disp_std
    return wrapping
user_129155310
Related