A simple in-memory Redis-like data store

simpledb.py

"""
Author: Mohamed Hendawy
Last edited: Feb 24 2015
Description: A simple in-memory Redis-like data store

HOW-TO: $ /path/to/python /path/to/simpledb.py < input.file > output.file

To Consider: The output alwas ends with enforced newline, but sometimes text
editors also add a newline to the end of the file
"""


class SimpleDB(object):
    """
    Core datastore, This class contains the basic methods and data
    structure for the datastore
    """

    def __init__(self):
        """
        Initializing the data structures required
        """
        self.__data_store = {}
        self.__value_count = {}

    def __setitem__(self, key, value):
        """
        Seting key to value
        """
        if key in self.__data_store \
                and self.__data_store[key] in self.__value_count \
                and self.__value_count[self.__data_store[key]] > 0:
            self.__value_count[self.__data_store[key]] -= 1
        if key in self.__data_store and \
                self.__data_store[key] in self.__value_count and \
                self.__value_count[self.__data_store[key]] < 1:
            del self.__value_count[self.__data_store[key]]
        if value in self.__value_count:
            self.__value_count[value] += 1
        else:
            self.__value_count[value] = 1
        self.__data_store[key] = value

    def __getitem__(self, key):
        """
        Getting value of key
        """
        if key in self.__data_store:
            return self.__data_store[key]
        else:
            return 'NULL'

    def __contains__(self, key):
        """
        Whether is a key in the datastore or not
        """
        return key in self.__data_store

    def __delitem__(self, key):
        """
        Unsetting a key
        """
        if key in self.__data_store \
                and self.__data_store[key] in self.__value_count \
                and self.__value_count[self.__data_store[key]] > 0:
            self.__value_count[self.__data_store[key]] -= 1
        if self.__data_store[key] in self.__value_count and \
                self.__value_count[self.__data_store[key]] < 1:
            del self.__value_count[self.__data_store[key]]
        del self.__data_store[key]

    def count_value(self, value):
        """
        returning occurences of a value in the datastore
        """
        if value in self.__value_count:
            return self.__value_count[value]
        else:
            return '0'


class SimpleDBCompiler(object):
    """
    Responsible for carrying on commands and keeping track of transactions
    """

    def __init__(self):
        """
        Initializing data structures and required objects
        """
        self.__transaction_stack = []
        self.__db = SimpleDB()
        self.__transaction_block = False
        self.action = None

    def __rollback(self):
        """
        Rolling back a transaction block
        """
        recent_stack = self.__transaction_stack.pop()
        while len(recent_stack) > 0:
            transaction = recent_stack.pop()
            transaction[0](*transaction[1])

    def compile(self, input_string):
        """
        Transforming raw commands into actions
        """
        partitions = input_string.split(' ')
        command, arg1, arg2 = None, None, None
        if len(partitions) > 0:
            command = partitions[0]
        if len(partitions) > 1:
            arg1 = partitions[1]
        if len(partitions) > 2:
            arg2 = partitions[2]

        if len(partitions) == 1:
            if command == 'BEGIN':
                self.__transaction_block = True
                self.__transaction_stack.append([])
            if command == 'COMMIT':
                self.__transaction_block = False
                self.__transaction_stack = []
            if command == 'ROLLBACK':
                if len(self.__transaction_stack) < 1 \
                        or len(self.__transaction_stack[-1]) < 1 \
                        or not self.__transaction_block:
                    return 'NO TRANSACTION'
                else:
                    self.__rollback()
            if command == 'END':
                self.action = exit
        elif len(partitions) > 1:
            if command == 'GET':
                return self.__db[arg1]
            elif command == 'NUMEQUALTO':
                return self.__db.count_value(arg1)
            elif command == 'SET':
                if self.__transaction_block:
                    if arg1 in self.__db:
                        self.__transaction_stack[-1].append(
                            (self.__db.__setitem__, (arg1, self.__db[arg1])))
                    else:
                        self.__transaction_stack[-1].append(
                            (self.__db.__delitem__, (arg1)))
                self.__db[arg1] = arg2
            elif command == 'UNSET':
                if self.__transaction_block:
                    if arg1 in self.__db:
                        self.__transaction_stack[-1].append(
                            (self.__db.__setitem__, (arg1, self.__db[arg1])))
                del self.__db[arg1]
        return ''


if __name__ == '__main__':
    sdbc = SimpleDBCompiler()
    while True:
        try:
            print sdbc.compile(raw_input())
            if sdbc.action is not None:
                sdbc.action()
        except EOFError:  # to handle EOF
            break