Python DB-Table Ish Thing

table.py

from collections import defaultdict


class Table:
    def __init__(self):
        self.items = []
        self.indices = {}

    @staticmethod
    def _get_key(item, key):
        val = None
        try:
            val = getattr(item, key)
        except AttributeError:
            try:
                val = item[key]
            except (TypeError, IndexError, KeyError):
                pass
        return val

    def index(self, key):
        if key in self.indices:
            return
        self.indices[key] = defaultdict(list)
        for item in self.items:
            self.index_item(item, key)

    def remove(self, item):
        self.items.remove(item)
        for index in self.indices:
            val = Table._get_key(item, index)
            self.indices[index][val].remove(item)

    def add(self, item):
        self.items.append(item)
        for key in self.indices:
            self.index_item(item, key)

    def index_item(self, item, key):
        val = Table._get_key(item, key)
        self.indices[key][val].append(item)

    def get(self, **kwargs):
        # @Optimization: do query+intersection over object ids or something?
        key = next(iter(kwargs))
        for key in self.indices:
            if key not in self.indices:
                raise KeyError("{} is not an indexed key".format(key))
        items = self.indices[key].get(kwargs[key], [])
        return [item for item in items
                if all(Table._get_key(item, k) == v
                       for k, v in kwargs.items())]

    def iter(self):
        return iter(self.items)

test_table.py

from table import Table


def people_table():
    tbl = Table()
    tbl.add({'name': 'a', 'pronoun': 'she'})
    tbl.add({'name': 'b', 'pronoun': 'he'})
    tbl.add({'name': 'c', 'pronoun': 'they'})
    tbl.add({'name': 'd', 'pronoun': 'she'})
    tbl.add({'name': 'd'})
    tbl.add({'name': 'e', 'pronoun': 'they'})
    return tbl


def test_create():
    people_table()


def test_index():
    people = people_table()
    people.index('name')


def test_index_partial():
    people = people_table()
    people.index('pronoun')


def test_get():
    people = people_table()
    people.index('name')
    assert people.get(name='c') == [
        {'name': 'c', 'pronoun': 'they'},
    ]


def test_partial_get():
    people = people_table()
    people.index('pronoun')
    assert people.get(pronoun='they') == [
        {'name': 'c', 'pronoun': 'they'},
        {'name': 'e', 'pronoun': 'they'},
    ]
    assert people.get(pronoun=None) == [{'name': 'd'}]


def test_multi_query():
    people = people_table()
    people.index('pronoun')
    people.index('name')
    assert people.get(pronoun='she', name='d') == [
        {'name': 'd', 'pronoun': 'she'},
    ]


def test_remove():
    people = people_table()
    people.index('name')
    people.remove({'name': 'd'})