python function to re order execution count of jupyter nb cells

reorder_execution_count.py

import json


def reorder_execution_count(infile, outfile):
    with open(infile, encoding='utf-8') as f:
        nb = json.load(f)

    count = 1

    for c in nb['cells']:
        if 'execution_count' in c:
            c['execution_count'] = count
            
            if 'outputs' in c:
                for o in c['outputs']:
                    if 'execution_count' in o:
                        o['execution_count'] = count
            count += 1

    with open(outfile, 'w', encoding='utf-8') as outfile:
        json.dump(nb, outfile, indent=2)

Not entirely sure what expected behavior should be for multiple output cells from 1 input cell as I have not run into that scenario. For 1 output cell the expected_count should be the same as the input cell. I made the assumption that this would also be true for more than 1 ouput cell.