Create stratified k-folds from an indexed dataframe

def stratified_kfold(df, n_folds):
    """
    Create stratified k-folds from an indexed dataframe
    """
    sessions = pd.DataFrame.from_records(list(df.index.unique())).groupby(0).apply(lambda x: x[1].unique())
    sessions.apply(lambda x: np.random.shuffle(x))
    folds = []
    for i in range(n_folds):
        idx = sessions.apply(lambda x: pd.Series(x[i * (len(x) / n_folds):(i + 1) * (len(x) / n_folds)]))
        idx = pd.DataFrame(idx.stack().reset_index(level=1, drop=True)).set_index(0, append=True).index.values
        folds.append(df.loc[idx])
    return folds