def assign_colors(data, aes, gg, aes_name='color'):
    """
    Assigns colors to the given data based on the aes and adds the right legend

    We need to take a value an convert it into colors that we can actually
    plot. This means checking to see if we're colorizing a discrete or
    continuous value, checking if their is a colormap, etc.

    Parameters
    ----------
    data : DataFrame
        dataframe which should have colors assigned to
    aes : aesthetic
        mapping, including a mapping from color to variable
    gg : ggplot object, which holds information and gets a legend assigned
    aes_name : string
        Which aesthetic needs a color mapping; 'color' | 'fill'

    Returns
    -------
    data : DataFrame
        the changed dataframe
    legend_entry : dict
        An entry into the legend dictionary.
        Documented in `components.legend`
    """
    legend_entry = dict()
    if aes_name in aes:
        color_col = aes[aes_name]
        labels, scale_type, indices = get_labels(data, color_col)
        if scale_type == "continuous" :
            data, legend_entry = assign_continuous_colors(
                data, gg, aes_name, color_col, labels, indices)
        elif scale_type == "discrete" :
            data, legend_entry = assign_discrete_colors(
                data, gg, aes_name, color_col, labels)
        else :
            raise GgplotError("Unknow scale_type: '%s'" % scale_type)

    return data, legend_entry
dataset/ETHPy150Open yhat/ggplot/ggplot/components/colors.py/assign_colors/original