From cfa13fbe457b7bfbc6132035d78754b1402441a2 Mon Sep 17 00:00:00 2001 From: Cameron Blocker Date: Sat, 1 May 2021 22:01:50 -0400 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Florian Rathgeber --- nbstripout/_utils.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/nbstripout/_utils.py b/nbstripout/_utils.py index b2db434..79a87a3 100644 --- a/nbstripout/_utils.py +++ b/nbstripout/_utils.py @@ -45,7 +45,7 @@ def get_size(item): if isinstance(item, str): return len(item) elif isinstance(item, list): - return functools.reduce(lambda a, b: a + get_size(b), item, 0) + return sum(get_size(elem) for elem in item) elif isinstance(item, dict): return get_size(list(item.values())) else: @@ -121,12 +121,7 @@ def strip_output(nb, keep_output, keep_count, extra_keys='', max_size=0): cell['outputs'] = [] else: # otherwise remove outputs larger than max_size - for idx, output in enumerate(cell['outputs']): - if get_size(output) > max_size: - cell['outputs'][idx] = None - elif 'execution_count' in output and not keep_count: - output['execution_count'] = None - cell['outputs'] = [out for out in cell['outputs'] if out is not None] + cell['outputs'] = [output for output in cell['outputs'] if get_size(output) <= max_size] # If keep_output_this_cell, but not keep_count, strip the counts # from the output.