Fix key error problem and add tests #11

src/underscore.py

@@ -399,7 +399,10 @@ def e(value, index, *args):
         _.each(obj, e)
 
         if len(ns.result) == 1:
-            return ns.result[0]
+            try:
+                return ns.result[0]
+            except KeyError:
+                return list(ns.result.values())[0]
         return ns.result
 
     def groupBy(self, val):

tests/test_collections.py

@@ -149,6 +149,8 @@ def test_groupby(self):
         self.assertTrue(0 in parity and 1 in parity, 'created a group for each value')
         self.assertEqual(_(parity[0]).join(', '), '2, 4, 6', 'put each even number in the right group')
 
+        self.assertEqual(_.groupBy([1], lambda num, *args: num), [1])
+
         llist = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"]
         grouped = _.groupBy(llist, lambda x, *args: len(x))
         self.assertEqual(_(grouped[3]).join(' '), 'one two six ten')
@@ -160,6 +162,8 @@ def test_countby(self):
         self.assertEqual(parity[True], 2)
         self.assertEqual(parity[False], 3)
 
+        self.assertEqual(_.countBy([1], lambda num, *args: num), 1)
+
         llist = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"]
         grouped = _.countBy(llist, lambda x, *args: len(x))
         self.assertEqual(grouped[3], 4)
@@ -216,6 +220,8 @@ def test_indexBy(self):
         self.assertEqual(parity[True], 4)
         self.assertEqual(parity[False], 5)
 
+        self.assertEqual(_.indexBy([1], lambda num, *args: num), 1)
+
         llist = ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"]
         grouped = _.indexBy(llist, lambda x, *args: len(x))
         self.assertEqual(grouped[3], 'ten')