Refactored a bit

Author: calebmadrigal

Diff for: hashtable.py

@@ -16,53 +16,59 @@ def __init__(self, capacity=1000):
         # The outmost list is the one which the hash function maps the index to. The next inner
         # Array is the list of objects in that storage cell. The 3rd level is the individual
         # item array, where the 1st item is the key, and the 2nd item is the value.
-        self.data = [[] for i in range(capacity)]
+        self.data = [[] for _ in range(capacity)]
+
+    def _find_by_key(self, key, find_result_func):
+        index = hash_function(key, self.capacity)
+        hash_table_cell = self.data[index]
+        found_item = None
+        for item in hash_table_cell:
+            if item[0] == key:
+                found_item = item
+                break
+
+        return find_result_func(found_item, hash_table_cell)
 
     def set(self, key, obj):
-        """ Insert object with key into hashtable. If key already exists, then the object will be
+        """ Insert object with key into hash table. If key already exists, then the object will be
         updated. Key must be a string. Returns self. """
 
-        index = hash_function(key, self.capacity)
-        storage_cell = self.data[index]
-        for item in storage_cell:
-            if item[0] == key:
-                item[1] = obj
-        else:
-            storage_cell.append([key, obj])
-            self.size += 1
+        def find_result_func(found_item, hash_table_cell):
+            if found_item:
+                found_item[1] = obj
+            else:
+                hash_table_cell.append([key, obj])
+                self.size += 1
 
+        self._find_by_key(key, find_result_func)
         return self
 
     def get(self, key):
         """ Get object with key (key must be a string). If not found, it will raise a KeyError. """
 
-        index = hash_function(key, self.capacity)
-        storage_cell = self.data[index]
-        for item in storage_cell:
-            if item[0] == key:
-                return item[1]
-        else:
-            raise KeyError(key)
+        def find_result_func(found_item, _):
+            if found_item:
+                return found_item[1]
+            else:
+                raise KeyError(key)
+
+        return self._find_by_key(key, find_result_func)
 
     def remove(self, key):
         """ Remove the object associated with key from the hashtable. If found, the object will
         be returned. If not found, KeyError will be raised. """
 
-        index = hash_function(key, self.capacity)
-        storage_cell = self.data[index]
-        item_to_remove = None
-        for item in storage_cell:
-            if item[0] == key:
-                item_to_remove = item
-        if item_to_remove:
-                storage_cell.remove(item)
+        def find_result_func(found_item, hash_table_cell):
+            if found_item:
+                hash_table_cell.remove(found_item)
                 self.size -= 1
-                return item[1]
-        else:
-            raise KeyError(key)
+                return found_item[1]
+            else:
+                raise KeyError(key)
+
+        return self._find_by_key(key, find_result_func)
 
 if __name__ == "__main__":
     import unittest
     testsuite = unittest.TestLoader().discover('test', pattern="*hashtable*")
     unittest.TextTestRunner(verbosity=1).run(testsuite)
-

Diff for: test/test_hashtable.py

@@ -32,23 +32,27 @@ def testCreate(self):
         self.assertEqual(ht.capacity, 1000)
 
     def testSetAndGet(self):
+        # Basic set and get
         ht = hashtable.HashTable(10)
         ht.set('a', 1)
         self.assertEqual(ht.get('a'), 1)
         self.assertEqual(ht.size, 1)
 
+        # Check update functionality
         ht.set('a', 2)
         self.assertEqual(ht.get('a'), 2)
-        #self.assertEqual(ht.size, 1)
+        self.assertEqual(ht.size, 1)
 
+        # Make sure we can add a 2nd element
         ht.set('b', 10)
         self.assertEqual(ht.get('b'), 10)
         self.assertEqual(ht.get('a'), 2)
-        #self.assertEqual(ht.size, 2)
+        self.assertEqual(ht.size, 2)
 
         # Assert ht.set returns itself (for fluent calls)
         self.assertEqual(ht.set('c', 5), ht)
 
+        # Test fluent set functionality
         ht.set('d', 100).set('e', 200).set('f', 300)
         self.assertEqual(ht.get('d'), 100)
         self.assertEqual(ht.get('e'), 200)
@@ -67,7 +71,6 @@ def testRemove(self):
         self.assertEqual(removed_item, 1)
         self.assertEqual(ht.size, 0)
 
-        
 if __name__ == '__main__':
     unittest.main()