data_handler.py 1.52 KB
Newer Older
1 2 3 4 5 6 7
"""
General IO for pickle database operations.
"""

import pickle


Kyle Anderson's avatar
Kyle Anderson committed
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
def get_user_ids_in_database(database: dict) -> list:
    """
    Gets all the user_ids in the given database.
    :param database: The database to look through.
    :return: All the user_ids in the database.
    """
    return list(database.keys())


def get_encodings_in_database(database: dict):
    """
    Gets a list of all encodings in the given database.
    :param database: The database dictionary.
    :return: All the encodings
    """
    result = []
    for encodings in database.values():
        result.extend(encodings)
    return result


29 30 31 32 33 34
def load_database(file_location: str):
    """
    Attempts to load the pickle database at the given file location
    :param file_location: String location of file to be loaded.
    :return: The loaded pickle database.
    """
Kyle Anderson's avatar
Kyle Anderson committed
35 36 37 38 39 40
    file_content = {}
    try:
        with open(file_location, "rb") as database:
            file_content = pickle.load(database)
    except (FileNotFoundError, EOFError):
        file_content = {}
41 42 43
    return file_content


Kyle Anderson's avatar
Kyle Anderson committed
44
def write_database(output_file: str, database_content) -> None:
45 46 47 48 49 50
    """
    Writes the dictionary database to the given file location
    :param output_file: The location of the file to be outputted on.
    :param database_content: The database content to be written to the file.
    :return: None
    """
Kyle Anderson's avatar
Kyle Anderson committed
51 52
    if output_file and database_content and database_content is not None:
        with open(output_file, "wb") as output:
Kyle Anderson's avatar
Kyle Anderson committed
53
            pickle.dump(database_content, output)