python-voting/voting/storage.py

import shutil

from arrow import get as arrow_get
from pathlib import Path
from simple_term_menu import TerminalMenu
from typing import List

from voting import Voting


class Storage(object):
    """Storage serves as a storage interface.

    It automatically loads `.voting.db` when initialized and creates the file if necessary.
    The persisted votings can then be access via the `votings` attribute.

    When used with the `with` statement, changes will automatically be saved when leaving the context.
    """
    def __init__(self, path: str = ".voting.db"):
        self.path: str = path
        self.votings: List = []
        Path(self.path).touch()
        self.load()

        self.__dirty = False

    def __enter__(self):
        return self
    
    def __exit__(self, exection_type, exception_value, stacktrace):
        if not exection_type:
            self.save()

    def load(self) -> None:
        """Loads the persisted votings from disk."""
        self.votings = []
        with open(Path(self.path), 'r') as f:
            for line in f.readlines():
                if line:
                    self.votings.append(Voting.loads(line))

    def save(self) -> None:
        """Saves the votings in memory to disk.
        
        The file is only ever overwritten if changes have actually occured.
        A backup file is created before saving called `.voting.db.bak`.
        """
        if not self.__dirty:
            self.__dirty = False
            return

        shutil.copyfile(".voting.db", ".voting.db.bak", follow_symlinks=True)  # kein backup: kein mitleid!
        lines = []
        for voting in self.votings:
            lines.append(voting.dumps())

        with open(Path(self.path), 'w') as f:
            for line in lines:
                f.write(f'{line}\n')

    def push(self, voting: Voting) -> None:
        """Push a voting to the in-memory voting list."""
        self.votings.append(voting)
        self.__dirty = True

    def pop(self, title: str) -> Voting:
        """Pop a voting from the in-memory voting list.
        
        If you call `save` after popping, the popped voting will effectively have been deleted from disk.
        In order to update a voting, firsr pop it, then manipulate it and push it back to the storage before calling `save`.

        In case the title is ambiguous, an interactive ASCII menu will be rendered for selection. 
        """
        candidates = []
        for v in self.votings:
            if v.title == title:
                candidates.append(v)
        
        if len(candidates) == 0:
            return None

        selection = 0
        if len(candidates) > 1:
            menu = TerminalMenu(
                [
                    f"[{i}] - {c.title} ({arrow_get(c.start).format('YYYY-MM-DD HH:mm:ss')})"
                    for i, c in enumerate(candidates)
                ],
                title="Candidates"
            )
            selection = menu.show()

        if selection is None:
            return None

        counter = 0
        votings = []
        popped = None
        for v in self.votings:
            if v.title == title:
                if counter == selection:
                    popped = v
                    self.__dirty = True
                else:
                    votings.append(v)
                counter += 1
                continue
            votings.append(v)
        self.votings = votings
        return popped