hugohm/utils/csveditor.py

#!/usr/bin/env python3
from argparse import ArgumentParser
import atexit
import csv
import os
import readline
import sys
from cmd import Cmd

def init_readline():
    histfile = os.path.join(os.path.expanduser("~"), ".hmwebot_history")

    try:
        readline.read_history_file(histfile)
        h_len = readline.get_current_history_length()
    except FileNotFoundError:
        open(histfile, 'wb').close()
        h_len = 0

    def save(prev_h_len, histfile):
        new_h_len = readline.get_current_history_length()
        readline.set_history_length(1000)
        readline.append_history_file(new_h_len - prev_h_len, histfile)
    atexit.register(save, h_len, histfile)


def rlinput(prompt, prefill=''):
   readline.set_startup_hook(lambda: readline.insert_text(prefill))
   try:
      return input(prompt)  # or raw_input in Python 2
   finally:
      readline.set_startup_hook()


class CSVEditor:
    def __init__(self, filename):
        self.filename = filename
        with open(filename, 'r') as f:
            reader = csv.DictReader(f)
            self.fieldnames = reader.fieldnames
            self.data = [data for data in reader]

    def get_list(self, title_field):
        return [(idx, data[title_field]) for idx, data in enumerate(self.data)]

    def get(self, idx):
        try:
            return self.data[idx]
        except IndexError:
            return None

    def show(self, idx):
        line = self.get(idx)
        if line is None:
            return "Índice fuera de rango"
        output = [
            f"{field}: {value}"
            for field, value in line.items()
        ]
        return "\n".join(output)

    def edit(self, idx, item):
        self.data[idx] = item

    def delete(self, idx):
        if idx in self.data:
            del self.data[idx]

    def _edit(self, item):
        editing = True
        while editing:
            for fieldname in self.fieldnames:
                item[fieldname] = rlinput(f"{fieldname}: ", prefill=item.get(fieldname, ''))

            editing = rlinput('Terminado (s/N)?').lower() == 'n'

    def edit(self, idx):
        item = self.get(idx)
        if item is None:
            return
        self._edit(item)
        self.data[idx] = item

    def add(self):
        item = {}
        self._edit(item)
        self.data.append(item)

    def write(self):
        with open(self.filename, 'w') as f:
            writer = csv.DictWriter(f, fieldnames=self.fieldnames)
            writer.writeheader()
            writer.writerows(self.data)


class CSVEditorCmd(Cmd):
    "Editor interactivo de ficheros csv"

    def __init__(self, filename, title_field=None):
        super().__init__()
        self.csveditor = CSVEditor(filename)
        self.title_field = title_field
        if title_field is None:
            self.title_field = self.csveditor.fieldnames[0]
        self.do_help('')

    def do_list(self, arg):
        "Listar el contenido del csv"
        for idx, item in self.csveditor.get_list(self.title_field):
            print(f"{idx}) {item}")

    def do_show(self, arg):
        "Mostrar la línea con el índice indicado"
        try:
            idx = int(arg)
        except ValueError:
            print("Tienes que indicar un número")
            return
        print(self.csveditor.show(idx))

    def do_add(self, arg):
        "Añadir una nueva entrada"
        self.csveditor.add()

    def do_edit(self, arg):
        "Editar la entrada con el índice indicado"
        try:
            idx = int(arg)
        except ValueError:
            print("Tienes que indicar un número")
            return
        self.csveditor.edit(idx)

    def do_delete(self, arg):
        "Quitar la entrada con el índice indicado"
        try:
            idx = int(arg)
        except ValueError:
            print("Tienes que indicar un número")
            return
        self.csveditor.delete(idx)

    def do_save(self, arg):
        "Guardar la lista actualizada"
        self.csveditor.write()

    def do_quit(self, arg):
        "Terminar"
        return True

    do_EOF = do_quit


parser = ArgumentParser(
    description='Editor sencillo de ficheros CSV')

parser.add_argument('csvfile', help='Fichero que hay que editar.')
parser.add_argument('-t', '--title', help='Columna principal del CSV, por defecto la primera.')


def main():
    init_readline()
    # bot = Bot()
    args = parser.parse_args()
    cmd = CSVEditorCmd(args.csvfile, args.title)
    cmd.cmdloop()

if __name__ == '__main__':
    main()