Skip to content
This repository has been archived by the owner on Jan 13, 2023. It is now read-only.

Commit

Permalink
Support adding/removing address checksums
Browse files Browse the repository at this point in the history
Two new methods in Address class:

 - add_checksum: appends a valid checksum to
   the Address object.

 - remove_checksum: slices off the checksum from
   an Address object.
  • Loading branch information
lzpap committed Nov 15, 2019
1 parent 86b53fa commit 6b544b4
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 0 deletions.
22 changes: 22 additions & 0 deletions iota/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -870,6 +870,28 @@ def _generate_checksum(self):

return AddressChecksum.from_trits(checksum_trits[-checksum_length:])

def add_checksum(self):
# type: () -> None
"""
Add checksum to :py:class:`Address` object.
"""
if self.is_checksum_valid():
# Address already has a valid checksum.
return

# Fill checksum attribute
self.checksum = self._generate_checksum()

# Add generated checksum to internal buffer.
self._trytes = self._trytes + self.checksum._trytes

def remove_checksum(self):
# type: () -> None
"""
Remove checksum from :py:class:`Address` object.
"""
self.checksum = None
self._trytes = self._trytes[:self.LEN]

class AddressChecksum(TryteString):
"""
Expand Down
33 changes: 33 additions & 0 deletions test/types_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1049,6 +1049,39 @@ def test_with_checksum_attributes(self):
self.assertEqual(checked.key_index, 42)
self.assertEqual(checked.balance, 86)

def test_add_checksum(self):
"""
Checksum is added to an address without it.
"""
addy = Address(
trytes =
b'ZKIUDZXQYQAWSHPKSAATJXPAQZPGYCDCQDRSMWWCGQJNI'
b'PCOORMDRNREDUDKBMUYENYTFVUNEWDBAKXMV'
)

addy.add_checksum()

self.assertTrue(addy.is_checksum_valid())
self.assertTrue(len(addy) == Address.LEN + AddressChecksum.LEN)

def test_remove_checksum(self):
"""
Checksum is removed from an address.
"""
addy = Address(
trytes =
b'ZKIUDZXQYQAWSHPKSAATJXPAQZPGYCDCQDRSMWWCGQJNI'
b'PCOORMDRNREDUDKBMUYENYTFVUNEWDBAKXMVJJJGBARPB'
)

self.assertTrue(addy.is_checksum_valid())
self.assertTrue(len(addy) == Address.LEN + AddressChecksum.LEN)

addy.remove_checksum()

self.assertFalse(addy.is_checksum_valid())
self.assertTrue(len(addy) == Address.LEN)


# noinspection SpellCheckingInspection
class AddressChecksumTestCase(TestCase):
Expand Down

0 comments on commit 6b544b4

Please sign in to comment.