From 6b544b4030a4dc5a4a0ef1719d1b400c1bdb7fa9 Mon Sep 17 00:00:00 2001
From: Levente Pap <levente.pap@iota.org>
Date: Fri, 15 Nov 2019 13:50:40 +0100
Subject: [PATCH] Support adding/removing address checksums

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.
---
 iota/types.py      | 22 ++++++++++++++++++++++
 test/types_test.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/iota/types.py b/iota/types.py
index 44fc58a..beeab33 100644
--- a/iota/types.py
+++ b/iota/types.py
@@ -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):
     """
diff --git a/test/types_test.py b/test/types_test.py
index 3ab0eb7..060db1f 100644
--- a/test/types_test.py
+++ b/test/types_test.py
@@ -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):