commit 3b46d49e499b9a15c1657bb643d0fcb7bd505faa Author: Damian Johnson <atagar@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'))