Rename cubes and coords in cubelist with unique names

rename_cubelist.py

import iris
import copy

# Main rename_cubelist() function
def rename_cubelist(cubelist, cubenames=None, new_coordnames=None, dryrun=False, verbose=True):
    '''Rename cubes and coordinates in place where necessary'''
    
    if cubenames==None:
        cubenames = [cube.name() for cube in cubelist]
        
    if new_coordnames==None:
        new_coordnames = get_new_coord_names(unique_coords_list(cubelist))
    
    for cube in cubelist:
        # Rename cube if duplicate or unknown
        if cube.standard_name == None or cubenames.count(cube.name()) > 1:
            new_name = get_new_cubename(cube)
            if dryrun or verbose:
                print(f'{cube.name()} -> {new_name}')
            if not dryrun:
                cube.var_name = new_name
        elif dryrun or verbose:
            print(f'{cube.name()}')
        
        # Rename coords
        for coord in cube.coords():
            if coord in new_coordnames[0]:
                new_name = new_coordnames[1][new_coordnames[0].index(coord)]
                if not dryrun:
                    coord.var_name = new_name
                if dryrun or verbose:
                    print(f'    {new_name}')
            elif dryrun or verbose:
                print(f'  x {coord.name()}')


# Additional functions used by rename_cubes()
def unique_coords_list(cubelist):
    '''Find all unique coords that represent every cube in the cubelist'''
    
    unique = []
    for cube in cubelist:
        for coord in cube.coords():
            if not coord in unique:
                unique.append(coord)
    return copy.deepcopy(unique)


def get_new_coord_names(coords, verbose=False):
    '''Rename coords that have the same name but different values'''
    
    names = []
    renamed = []
    for coord in coords:
        name = coord.name()
        names.append(name)
        n = names.count(name)
        if n > 1:
            new_name = f'{name}_{n-1}'
            renamed.append((coord, new_name))
        if verbose:
            print(f'Names: {names}')
    if verbose:
        print(f'Names: {names}')
    return tuple(zip(*renamed))


def get_new_cubename(cube):
    '''Rename cube based on its vertical coord and cell method'''
    
    suffixes = [cube.standard_name or str(cube.attributes['STASH'])]  
    # [cube.name()] leads to repeated cell_method suffixes for anonymous cubes
    coord_names = [coord.name() for coord in cube.coords()]
    
    if 'pressure' in coord_names:
        suffixes.append('at_pressure')
    
    if 'height' in coord_names:
        heights = cube.coord('height')
        if len(heights.points) > 1:
            suffixes.append('at_height')
        else:
            height = str(int(heights.points[0].round()))
            units = str(heights.units)
            suffixes.append(f'at_{height}{units}')
    
    for cell_method in cube.cell_methods:
        method = cell_method.method.replace('imum', '')
        suffixes.append(method)
    
    return '_'.join(suffixes)