WPE #15 CheckPickle

solution.py

from cmd import Cmd
from glob import glob
import pickle
import re
from time import time
from typing import Dict, List


fields = ["first_name", "last_name", "email"]

TAddressBook = List[Dict[str, str]]


class AddressBook(Cmd):
    prompt = "AdB> "

    def __init__(self, cp_stem: str):
        super().__init__()
        self.book = list()  # type: TAddressBook
        self.cp_stem = cp_stem
        self.aliases = {
            "a": self.do_add,
            "h": self.do_help,
            "l": self.do_list,
            "q": self.do_quit,
            "r": self.do_restore,
        }

    def do_list(self, inp):
        """List all people in the address book"""
        print("== Persons already recorded ==")
        for p in self.book:
            first = p["first_name"].title()
            last = p["last_name"].title()
            print(f" * {first} {last} ({p['email']})")

    def do_add(self, inp):
        """Add a new person to the address book"""
        person = dict()  # type: Dict[str, str]
        person["first_name"] = input("First name: ")
        person["last_name"] = input("Last name: ")
        person["email"] = input("Email: ")
        # We loop for email until valid (searching for '@' and '.')
        while not re.match(r"[\w\+\-\.]+@\w+\.\w+", person["email"]):
            print("... Not a valid email!")
            person["email"] = input("Email: ")
        # We record the list of persons in the object and dump into a pickle
        self.book.append(person)
        with open(f"{self.cp_stem}-{time()}", "wb") as f:
            pickle.dump(self.book, f)

    def do_restore(self, inp):
        """Restore the address book to the stage from a specific timestamp"""
        if inp:
            try:
                with open(inp, "rb") as f:
                    self.book = pickle.load(f)
                print(f"Restored record {inp!r}")
            except Exception as e:
                print(f"ERROR: {inp} is not a valid record: {e}")
        else:
            print("Nothing done - no record provided")

    def complete_restore(self, text: str, line, begidx, endidx):
        """Auto-completion to help finding available records"""
        list_pickles = glob(f"{self.cp_stem}*")
        if not text:
            completions = list_pickles
        else:
            completions = [f for f in list_pickles if f.startswith(text)]
        return completions

    def do_quit(self, inp):
        """Quit from the program (function)"""
        return True

    def do_help(self, arg):
        """List available commands."""
        if arg in self.aliases:
            arg = self.aliases[arg].__name__[3:]
        super().do_help(arg)

    def default(self, line):
        """Default, to allow aliases"""
        cmd, arg, line = self.parseline(line)
        if cmd in self.aliases:
            return self.aliases[cmd](arg)
        else:
            print("*** Unknown syntax: %s" % line)


def checkpickle(cp_stem="people-checkpoint-"):
    """Managing Address Book"""
    AddressBook(cp_stem).cmdloop()


if __name__ == "__main__":
    checkpickle()