Correct sizeof implementation (#12)

.travis.yml

@@ -9,7 +9,6 @@ python:
   - 2.7
   - 3.5
   - 3.6
-  - pypy
 
 install: false
 script:

src/kinesis/producer.py

@@ -30,21 +30,19 @@ def sizeof(obj, seen=None):
 
     size = sys.getsizeof(obj)
 
-    # since strings are containers we return their size before we check for a container
+    # since strings are iterabes we return their size explicitly first
     if isinstance(obj, six.string_types):
         return size
-
-    if isinstance(obj, collections.Container):
-        return size + sum(
-            sizeof(item, seen)
-            for item in obj
-        )
-
-    if isinstance(obj, collections.Mapping):
+    elif isinstance(obj, collections.Mapping):
         return size + sum(
             sizeof(key, seen) + sizeof(val, seen)
             for key, val in six.iteritems(obj)
         )
+    elif isinstance(obj, collections.Iterable):
+        return size + sum(
+            sizeof(item, seen)
+            for item in obj
+        )
 
     return size
 
@@ -116,6 +114,7 @@ def loop(self):
                 break
 
         self.flush_records()
+        return 0
 
     def end(self):
         # At the end of our loop (before we exit, i.e. via a signal) we change our buffer time to 250ms and then re-call

test/test_producer.py

@@ -1,4 +1,6 @@
-from kinesis.producer import KinesisProducer
+import sys
+
+from kinesis.producer import KinesisProducer, sizeof
 
 
 def test_producer(mocker):
@@ -16,3 +18,15 @@ def test_producer(mocker):
     mocked_queue = mocker.patch.object(producer, 'queue')
     producer.put('foo', explicit_hash_key='hash', partition_key='partition')
     mocked_queue.put.assert_called_with(('foo', 'hash', 'partition'))
+
+
+def test_sizeof():
+    s1 = "a"
+    assert sizeof(s1) == sys.getsizeof(s1)
+
+    s2 = ["a", "b", "a"]
+    assert sizeof(s2) == sys.getsizeof(s2) + sys.getsizeof("a") + sys.getsizeof("b")
+
+    s3 = {"a": "a", "b": "a", "c": 1}
+    assert sizeof(s3) == sys.getsizeof(s3) + sys.getsizeof("a") + sys.getsizeof("b") + sys.getsizeof("c") + \
+        sys.getsizeof(1)