[tor-commits] [stem/master] Drop struct usage from Size class

[atagar at torproject.org]

commit 3b46d49e499b9a15c1657bb643d0fcb7bd505faa
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Jan 4 17:44:48 2020 -0800

    Drop struct usage from Size class
    
    Python 3.x added better methods for integer conversion to/from bytes.
---
 stem/client/datatype.py  | 38 +++++++++++---------------------------
 test/unit/client/size.py |  6 ++----
 2 files changed, 13 insertions(+), 31 deletions(-)

diff --git a/stem/client/datatype.py b/stem/client/datatype.py
index 0b39d1b5..1db46a28 100644
--- a/stem/client/datatype.py
+++ b/stem/client/datatype.py
@@ -138,7 +138,6 @@ users.** See our :class:`~stem.client.Relay` the API you probably want.
 import binascii
 import collections
 import hashlib
-import struct
 
 import stem.client.cell
 import stem.prereq
@@ -367,7 +366,7 @@ class Field(object):
 class Size(Field):
   """
   Unsigned `struct.pack format
-  <https://docs.python.org/2/library/struct.html#format-characters>` for
+  <https://docs.python.org/3/library/struct.html#format-characters>` for
   network-order fields.
 
   ====================  ===========
@@ -380,45 +379,30 @@ class Size(Field):
   ====================  ===========
   """
 
-  def __init__(self, name, size, pack_format):
+  def __init__(self, name, size):
     self.name = name
     self.size = size
-    self.format = pack_format
 
   @staticmethod
   def pop(packed):
     raise NotImplementedError("Use our constant's unpack() and pop() instead")
 
   def pack(self, content):
-    # TODO: When we drop python 2.x support this can be simplified via
-    # integer's to_bytes() method. For example...
-    #
-    #   struct.pack('>Q', my_number)
-    #
-    # ... is the same as...
-    #
-    #   my_number.to_bytes(8, 'big')
-
     try:
-      packed = struct.pack(self.format, content)
-    except struct.error:
+      return content.to_bytes(self.size, 'big')
+    except:
       if not isinstance(content, int):
         raise ValueError('Size.pack encodes an integer, but was a %s' % type(content).__name__)
       elif content < 0:
         raise ValueError('Packed values must be positive (attempted to pack %i as a %s)' % (content, self.name))
       else:
-        raise  # some other struct exception
-
-    if self.size != len(packed):
-      raise ValueError('%s is the wrong size for a %s field' % (repr(packed), self.name))
-
-    return packed
+        raise
 
   def unpack(self, packed):
     if self.size != len(packed):
       raise ValueError('%s is the wrong size for a %s field' % (repr(packed), self.name))
 
-    return struct.unpack(self.format, packed)[0]
+    return int.from_bytes(packed, 'big')
 
   def pop(self, packed):
     to_unpack, remainder = split(packed, self.size)
@@ -426,7 +410,7 @@ class Size(Field):
     return self.unpack(to_unpack), remainder
 
   def __hash__(self):
-    return stem.util._hash_attr(self, 'name', 'size', 'format', cache = True)
+    return stem.util._hash_attr(self, 'name', 'size', cache = True)
 
 
 class Address(Field):
@@ -722,7 +706,7 @@ def _unpack_ipv6_address(value):
   return ':'.join(['%04x' % Size.SHORT.unpack(value[i * 2:(i + 1) * 2]) for i in range(8)])
 
 
-setattr(Size, 'CHAR', Size('CHAR', 1, '!B'))
-setattr(Size, 'SHORT', Size('SHORT', 2, '!H'))
-setattr(Size, 'LONG', Size('LONG', 4, '!L'))
-setattr(Size, 'LONG_LONG', Size('LONG_LONG', 8, '!Q'))
+setattr(Size, 'CHAR', Size('CHAR', 1))
+setattr(Size, 'SHORT', Size('SHORT', 2))
+setattr(Size, 'LONG', Size('LONG', 4))
+setattr(Size, 'LONG_LONG', Size('LONG_LONG', 8))
diff --git a/test/unit/client/size.py b/test/unit/client/size.py
index dcd5fea3..d3bb4a4a 100644
--- a/test/unit/client/size.py
+++ b/test/unit/client/size.py
@@ -11,8 +11,6 @@ from stem.client.datatype import Size
 class TestSize(unittest.TestCase):
   def test_attributes(self):
     self.assertEqual('CHAR', Size.CHAR.name)
-    self.assertEqual('!B', Size.CHAR.format)
-
     self.assertEqual(1, Size.CHAR.size)
     self.assertEqual(2, Size.SHORT.size)
     self.assertEqual(4, Size.LONG.size)
@@ -27,8 +25,8 @@ class TestSize(unittest.TestCase):
     self.assertRaisesWith(ValueError, 'Size.pack encodes an integer, but was a str', Size.CHAR.pack, 'hi')
     self.assertRaisesWith(ValueError, 'Packed values must be positive (attempted to pack -1 as a CHAR)', Size.CHAR.pack, -1)
 
-    bad_size = Size('BAD_SIZE', 1, '!H')
-    self.assertRaisesRegexp(ValueError, re.escape("'\\x00\\x12' is the wrong size for a BAD_SIZE field"), bad_size.pack, 18)
+    too_small = Size('TOO_SMALL', 1)
+    self.assertRaises(OverflowError, too_small.pack, 1800)
 
   def test_unpack(self):
     self.assertEqual(18, Size.CHAR.unpack(b'\x12'))