Merge pull request #112 from plenarius/chop_checksum_for_core_commands

[#101] Filter that automatically removes checksums for core commands

Author: todofixthis

iota/commands/core/find_transactions.py

@@ -7,7 +7,7 @@
 
 from iota import Address, Tag, TransactionHash
 from iota.commands import FilterCommand, RequestFilter, ResponseFilter
-from iota.filters import Trytes
+from iota.filters import AddressNoChecksum, Trytes
 
 __all__ = [
   'FindTransactionsCommand',
@@ -43,7 +43,7 @@ def __init__(self):
             f.Array
           | f.FilterRepeater(
                 f.Required
-              | Trytes(result_type=Address)
+              | AddressNoChecksum()
               | f.Unicode(encoding='ascii', normalize=False)
             )
         ),

iota/commands/core/get_balances.py

@@ -6,7 +6,7 @@
 
 from iota import Address
 from iota.commands import FilterCommand, RequestFilter, ResponseFilter
-from iota.filters import Trytes
+from iota.filters import AddressNoChecksum, Trytes
 
 __all__ = [
   'GetBalancesCommand',
@@ -37,7 +37,7 @@ def __init__(self):
           | f.Array
           | f.FilterRepeater(
                 f.Required
-              | Trytes(result_type=Address)
+              | AddressNoChecksum()
               | f.Unicode(encoding='ascii', normalize=False)
             )
         ),

iota/filters.py

@@ -142,3 +142,42 @@ def _apply(self, value):
           'result_type': self.result_type.__name__,
         },
       )
+
+
+class AddressNoChecksum(Trytes):
+  """
+  Validates a sequence as an Address then chops off the checksum if it exists
+  """
+  ADDRESS_BAD_CHECKSUM   = 'address_bad_checksum'
+
+  templates = {
+    ADDRESS_BAD_CHECKSUM: 'Checksum is {supplied_checksum}, should be {expected_checksum}?',
+  }
+
+  def __init__(self):
+    # type: (type) -> None
+    super(AddressNoChecksum, self).__init__(result_type=Address)
+
+  def _apply(self, value):
+    super(AddressNoChecksum, self)._apply(value)
+
+    if self._has_errors:
+      return None
+
+    # Possible it's still just a TryteString
+    if not isinstance(value, Address):
+      value = Address(value)
+
+    # Bail out if we have a bad checksum
+    if value.checksum and not value.is_checksum_valid():
+      return self._invalid_value(
+        value     = value,
+        reason    = self.ADDRESS_BAD_CHECKSUM,
+        exc_info  = True,
+
+        context = {
+          'supplied_checksum': value.checksum,
+          'expected_checksum': value.with_valid_checksum().checksum,
+        },
+      )
+    return Address(value.address)

test/filters_test.py

@@ -6,7 +6,7 @@
 from filters.test import BaseFilterTestCase
 
 from iota import Address, TryteString, TransactionHash
-from iota.filters import GeneratedAddress, NodeUri, Trytes
+from iota.filters import AddressNoChecksum, GeneratedAddress, NodeUri, Trytes
 
 
 class GeneratedAddressTestCase(BaseFilterTestCase):
@@ -201,3 +201,48 @@ def test_fail_wrong_type(self):
       [TryteString(b'RBTC9D9DCDQAEASBYBCCKBFA')],
       [f.Type.CODE_WRONG_TYPE],
     )
+
+
+# noinspection SpellCheckingInspection
+class AddressNoChecksumTestCase(BaseFilterTestCase):
+  filter_type = AddressNoChecksum
+
+  # noinspection SpellCheckingInspection
+  def setUp(self):
+    super(AddressNoChecksumTestCase, self).setUp()
+
+    # Define some addresses that we can reuse between tests
+    """
+    Incoming value is not an :py:class:`Address` instance.
+    """
+    self.tryte1 = (
+      b'TESTVALUE9DONTUSEINPRODUCTION99999FBFFTG'
+      b'QFWEHEL9KCAFXBJBXGE9HID9XCOHFIDABHDG9AHDR'
+    )
+    self.checksum = b'ENXYJOBP9'
+    self.address = Address(self.tryte1)
+    self.address_with_checksum = Address(self.tryte1 + self.checksum)
+    self.address_with_bad_checksum = Address(self.tryte1 + b'DEADBEEF9')
+
+  def test_pass_no_checksum_addy(self):
+    """
+    Incoming value is tryte in address form or Address object
+    """
+    self.assertFilterPasses(self.tryte1)
+    self.assertFilterPasses(self.address)
+
+  def test_pass_with_checksum_addy(self):
+    """
+    After passing through the filter an address with a checksum should
+    return the address without
+    """
+    self.assertFilterPasses(self.address_with_checksum, self.address)
+
+  def test_fail_with_bad_checksum_addy(self):
+    """
+    If they've got a bad checksum in their address we should probably tell
+    them so they don't wonder why something works in one place and not another
+    """
+    self.assertFilterErrors(
+      self.address_with_bad_checksum,
+      [AddressNoChecksum.ADDRESS_BAD_CHECKSUM])