From b6b5f697acba0a905d7021518cadccf41df9f222 Mon Sep 17 00:00:00 2001
From: Eli Bendersky <eliben@gmail.com>
Date: Wed, 30 Aug 2017 06:28:56 -0700
Subject: [PATCH] Clean up py3compat a bit and add some tests

---
 elftools/common/py3compat.py | 28 +++++++++++++++-------------
 test/test_py3compat.py       | 30 ++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 13 deletions(-)
 create mode 100644 test/test_py3compat.py

diff --git a/elftools/common/py3compat.py b/elftools/common/py3compat.py
index f481043..9285d66 100644
--- a/elftools/common/py3compat.py
+++ b/elftools/common/py3compat.py
@@ -1,7 +1,7 @@
 #-------------------------------------------------------------------------------
 # elftools: common/py3compat.py
 #
-# Python 3 compatibility code
+# Python 2/3 compatibility code
 #
 # Eli Bendersky (eliben@gmail.com)
 # This code is in the public domain
@@ -15,15 +15,22 @@ if PY3:
     StringIO = io.StringIO
     BytesIO = io.BytesIO
 
-    _iterkeys = "keys"
-    _iteritems = "items"
-    _itervalues = "values"
+    # Functions for acting on bytestrings and strings. In Python 2 and 3,
+    # strings and bytes are the same and chr/ord can be used to convert between
+    # numeric byte values and their string pepresentations. In Python 3, bytes
+    # and strings are different types and bytes hold numeric values when
+    # iterated over.
 
     def bytes2str(b): return b.decode('latin-1')
     def str2bytes(s): return s.encode('latin-1')
-    def int2byte(i):return bytes((i,))
+    def int2byte(i): return bytes((i,))
     def byte2int(b): return b
+
     def iterbytes(b):
+        """Return an iterator over the elements of a bytes object.
+        
+        For example, for b'abc' yields b'a', b'b' and then b'c'.
+        """
         for i in range(len(b)):
             yield b[i:i+1]
 
@@ -34,10 +41,6 @@ else:
     import cStringIO
     StringIO = BytesIO = cStringIO.StringIO
 
-    _iterkeys = "iterkeys"
-    _iteritems = "iteritems"
-    _itervalues = "itervalues"
-
     def bytes2str(b): return b
     def str2bytes(s): return s
     int2byte = chr
@@ -52,13 +55,12 @@ else:
 
 def iterkeys(d):
     """Return an iterator over the keys of a dictionary."""
-    return getattr(d, _iterkeys)()
+    return getattr(d, 'keys' if PY3 else 'iterkeys')()
 
 def itervalues(d):
     """Return an iterator over the values of a dictionary."""
-    return getattr(d, _itervalues)()
+    return getattr(d, 'values' if PY3 else 'itervalues')()
 
 def iteritems(d):
     """Return an iterator over the items of a dictionary."""
-    return getattr(d, _iteritems)()
-
+    return getattr(d, 'items' if PY3 else 'iteritems')()
diff --git a/test/test_py3compat.py b/test/test_py3compat.py
new file mode 100644
index 0000000..9c8fb13
--- /dev/null
+++ b/test/test_py3compat.py
@@ -0,0 +1,30 @@
+#-------------------------------------------------------------------------------
+# elftools tests
+#
+# Eli Bendersky (eliben@gmail.com)
+# This code is in the public domain
+#-------------------------------------------------------------------------------
+import unittest
+
+from elftools.common.py3compat import (iterbytes, iterkeys, itervalues,
+                                       iteritems)
+
+
+class TestPy3Compat(unittest.TestCase):
+    def test_iterbytes(self):
+        bi = iterbytes(b'fo1')
+        self.assertEqual(next(bi), b'f')
+        self.assertEqual(next(bi), b'o')
+        self.assertEqual(next(bi), b'1')
+        with self.assertRaises(StopIteration):
+            next(bi)
+
+    def test_iterdict(self):
+        d = {1: 'foo', 2: 'bar'}
+        self.assertEqual(list(sorted(iterkeys(d))), [1, 2])
+        self.assertEqual(list(sorted(itervalues(d))), ['bar', 'foo'])
+        self.assertEqual(list(sorted(iteritems(d))), [(1, 'foo'), (2, 'bar')])
+
+
+if __name__ == '__main__':
+    unittest.main()
-- 
2.30.2