Object Oriented Design - HashMap

In [11]:

### This is a simple HashMap class build in Jupyter Notebook
### It assumes only int key values
### It assumes there will not be any need to resize store size because of collision

In [12]:

class HashMap(object):
    
    class _Item(object):
        def __init__(self, key, value):
            self.key = key
            self.value = value
    
    def __init__(self, size):
        self.size = size
        self.store = [[] for _ in range(self.size)]
        
    def _hash(self, key):
        # assuming key is always int
        return key % self.size
    
    def put(self, key, value):
        index = self._hash(key)
        for item in self.store[index]:
            if item.key == key:
                item.value = value
                return
        self.store[index].append(self._Item(key, value))
    
    def get(self, key):
        index = self._hash(key)
        for item in self.store[index]:
            if item.key == key:
                return item.value
        raise KeyError('Key not found')
    
    def remove(self, key):
        index = self._hash(key)
        for item_index, item in enumerate(self.store[index]):
            if item.key == key:
                del self.store[index][item_index]
                return
        raise KeyError('Key not found')

In [13]:

import unittest

class HashMapTest(unittest.TestCase):
    
    def test_put_get(self):
        hm = HashMap(5)
        hm.put(1,'one')
        self.assertEqual(hm.get(1), 'one')
        hm.put(1,'updated_one')
        self.assertEqual(hm.get(1), 'updated_one')
        
    def test_remove(self):
        hm = HashMap(5)
        hm.put(1,'one')
        hm.remove(1)
        with self.assertRaises(KeyError):
            hm.get(1)
        
if __name__ == '__main__':
    unittest.main(argv=['first-arg-is-ignored'], exit=False)

..
----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK