Allow to manage a dataframe stored in a file (via read_pickle / write_pickle) from command line

df_manager.py

# ---------------------------------------------------------------------
# Allow to manage a dataframe stored in a file (via read_pickle / write_pickle)
# from command line
# ---------------------------------------------------------------------


import argparse
from collections import OrderedDict

import pandas as pd
import numpy as np


# ------ Actions funcs ------

def add_column(df):
    new_col = None
    while not new_col:
        new_col = raw_input('New column name: ').strip()
    if raw_input('Create column {} ? Y/N\n> '.format(new_col)).strip().lower() != 'y':
        print('Abort')
        return
    df[new_col] = pd.Series(np.nan, index=df.index)
    print('Success !')


def remove_column(df):
    col_name = _input_column_name(df, 'Column to remove: ')
    if raw_input('Delete column {} ? Y/N\n> '.format(col_name)).strip().lower() != 'y':
        print('Abort')
        return
    del df[col_name]
    print('Success !')


def rename_column(df):
    old_name = _input_column_name(df, 'Old column name: ')
    new_name = None
    while not new_name:
        new_name = raw_input('New column name: ').strip()
    if raw_input('Rename {} to {} ? Y/N\n> '.format(old_name, new_name)).strip().lower() != 'y':
        print('Abort')
        return
    df.columns = [col_name if col_name != old_name else new_name for col_name in df.columns]
    print('Success !')


def show_columns(df):
    print('{} columns : {}\n'.format(len(df.columns), ', '.join(df.columns)))


def leave(df=None):
    print('Bye !')

# ------------------------


ACTIONS = OrderedDict([
    ('1', add_column),
    ('2', remove_column),
    ('3', rename_column),
    ('4', show_columns),
    ('5', leave)
])

ACTIONS_STR = '\n'.join([' {}. {}'.format(shortcut, action.__name__) for shortcut, action in ACTIONS.items()])

MAIN_MENU_PROMPT = 'What do you want to do ?\n{}\n> '.format(ACTIONS_STR)


def _input_column_name(df, prompt):
    col_name = None
    while not col_name or col_name not in df.columns:
        col_name = raw_input(prompt).strip()
        if col_name not in df.columns:
            print('Invalid column name : {}. Columns are : {}'.format(col_name, ', '.join(df.columns)))
    return col_name


def main():
    # Commands params
    parser = argparse.ArgumentParser(description='Get benchmark request RPD status from API and generate a HTML page as result')
    parser.add_argument('-d', '--dataframe_path', dest='df_file', help="Path to the dataframe dump", required=True)
    cmd_options = parser.parse_args()

    df = pd.read_pickle(cmd_options.df_file)
    print('Dataframe loaded with {} columns : {}\n'.format(len(df.columns), ', '.join(df.columns)))
    print('----------------------------------------\n')
    user_input = ''
    while user_input != '5':
        user_input = raw_input(MAIN_MENU_PROMPT).strip()
        if user_input in ACTIONS:
            ACTIONS[user_input](df)
        else:
            print('Unknown command: {}'.format(user_input))

if __name__ == '__main__':
    main()