Add and fix witness_type in Address and Output class

bitcoinlib/keys.py

@@ -852,6 +852,10 @@ def __init__(self, data='', hashed_data='', prefix=None, script_type=None,
             elif self.script_type == 'p2tr':
                 witness_type = 'taproot'
                 self.witver = 1 if self.witver == 0 else self.witver
+            elif self.encoding == 'base58':
+                witness_type = 'legacy'
+            else:
+                witness_type = 'segwit'
         self.witness_type = witness_type
         self.depth = depth
         self.change = change

bitcoinlib/transactions.py

@@ -589,7 +589,7 @@ class Output(object):
 
     def __init__(self, value, address='', public_hash=b'', public_key=b'', lock_script=b'', spent=False,
                  output_n=0, script_type=None, witver=0, encoding=None, spending_txid='', spending_index_n=None,
-                 strict=True, change=None, network=DEFAULT_NETWORK):
+                 strict=True, change=None, witness_type=None, network=DEFAULT_NETWORK):
         """
         Create a new transaction output
         
@@ -628,6 +628,8 @@ def __init__(self, value, address='', public_hash=b'', public_key=b'', lock_scri
         :type strict: bool
         :param change: Is this a change output back to own wallet or not? Used for replace-by-fee.
         :type change: bool
+        :param witness_type: Specify witness type: 'segwit' or 'legacy'. Determine from script, address or encoding if not specified.
+        :type witness_type: str
         :param network: Network, leave empty for default
         :type network: str, Network
         """
@@ -652,42 +654,41 @@ def __init__(self, value, address='', public_hash=b'', public_key=b'', lock_scri
             public_key = address.public_byte
             if not script_type:
                 script_type = script_type_default(address.witness_type, address.multisig, True)
-            self.public_hash = address.hash160
+            # self.public_hash = address.hash160
+            # self.witness_type = address.witness_type
         else:
             self._address = address
             self._address_obj = None
         self.public_key = to_bytes(public_key)
         self.compressed = True
-        self.k = None
         self.versionbyte = self.network.prefix_address
         self.script_type = script_type
         self.encoding = encoding
-        if not self._address and self.encoding is None:
-            self.encoding = 'base58'
         self.spent = spent
         self.output_n = output_n
         self.script = Script.parse_bytes(self.lock_script, strict=strict, is_locking=True)
         self.witver = witver
+        self.witness_type = witness_type
 
         if self._address_obj:
             self.script_type = self._address_obj.script_type if script_type is None else script_type
+            # if not script_type:
+            #     script_type = script_type_default(address.witness_type, address.multisig, True)
             self.public_hash = self._address_obj.hash_bytes
             self.network = self._address_obj.network
             self.encoding = self._address_obj.encoding
+            self.witness_type = self._address_obj.witness_type
 
         if self.script:
             self.script_type = self.script_type if not self.script.script_types else self.script.script_types[0]
-            if self.script_type in ['p2wpkh', 'p2wsh', 'p2tr']:
-                self.encoding = 'bech32'
             self.public_hash = self.script.public_hash
             if self.script.keys:
                 self.public_key = self.script.keys[0].public_byte
             if self.script_type == 'p2tr':
                 self.witver = self.script.commands[0] - 80
 
         if self.public_key and not self.public_hash:
-            k = Key(self.public_key, is_private=False, network=network)
-            self.public_hash = k.hash160
+            self.public_hash = hash160(self.public_key)
         elif self._address and (not self.public_hash or not self.script_type or not self.encoding):
             address_dict = deserialize_address(self._address, self.encoding, self.network.name)
             if address_dict['script_type'] and not script_type:
@@ -703,10 +704,13 @@ def __init__(self, value, address='', public_hash=b'', public_key=b'', lock_scri
                 raise TransactionError("Network for output address %s is different from transaction network. %s not "
                                        "in %s" % (self._address, self.network.name, network_guesses))
             self.public_hash = address_dict['public_key_hash_bytes']
+            self.witness_type = address_dict['witness_type']
         if not self.encoding:
-            self.encoding = 'base58'
-            if self.script_type in ['p2wpkh', 'p2wsh', 'p2tr']:
-                self.encoding = 'bech32'
+            self.encoding = 'bech32'
+            if self.script_type in ['p2pkh', 'p2sh', 'p2pk'] or self.witness_type == 'legacy':
+                self.encoding = 'base58'
+            else:
+                self.witness_type = 'segwit'
 
         if self.script_type is None:
             self.script_type = 'p2pkh'

tests/test_keys.py

@@ -792,7 +792,7 @@ def test_hdkey_derive_from_public_and_private_random(self):
             self.assertEqual(pub_with_pubparent, pub_with_privparent)
 
 
-class TestKeysAddress(unittest.TestCase):
+class TestAddress(unittest.TestCase):
     """
     Tests for Address class. Address format, conversion and representation
 
@@ -898,6 +898,18 @@ def test_keys_address_p2tr_bcrt(self):
                        encoding='bech32').address
         self.assertEqual(addr, 'bcrt1pq77c6jeemv8wxlsh5h5pfdq6323naua8yapte3juw9hyec83mr8sw2eggg')
 
+    def test_keys_address_witness_types(self):
+        data = b'\x03\xb0\x12\x86\x15bt\xc9\x0f\xa7\xd0\xf6\xe6\x17\xc9\xc6\xafS\xa0u/ou\x8d\xa5\x1d\x1c\xc9h4nl\xb8'
+        a = Address(data)
+        self.assertEqual(a.address, 'bc1q36cn4tunsaptdskkf29lerzym0uznqw26pxffm')
+        self.assertEqual(a.witness_type, 'segwit')
+        a = Address(data, witness_type='segwit')
+        self.assertEqual(a.address, 'bc1q36cn4tunsaptdskkf29lerzym0uznqw26pxffm')
+        self.assertEqual(a.witness_type, 'segwit')
+        a = Address(data, witness_type='legacy')
+        self.assertEqual(a.address, '1E1VGLvZ2YpgcSgr3DYm7ZTHbovKw9xLw6')
+        self.assertEqual(a.witness_type, 'legacy')
+
 
 class TestKeysSignatures(unittest.TestCase):
 

tests/test_transactions.py

@@ -108,29 +108,37 @@ def test_transaction_input_with_pkh(self):
 
 class TestTransactionOutputs(unittest.TestCase):
 
-    def test_transaction_output_add_address(self):
+    def test_transaction_output_address(self):
         to = Output(1000, '1QhnmvncrbZFkjt5R8hs8yHDM7xXX3feg')
         self.assertEqual(b'v\xa9\x14\x04{\x9d\xc2=\xda\xa9\x17\x1e\xa5\x11\xe1\x93t\xabUo\xaa\xbbD\x88\xac',
                          to.lock_script)
         self.assertEqual(repr(to), '<Output(value=1000, address=1QhnmvncrbZFkjt5R8hs8yHDM7xXX3feg, type=p2pkh)>')
 
-    def test_transaction_output_add_address_p2sh(self):
+    def test_transaction_output_address_p2sh(self):
         to = Output(1000, '2N5WPJ2qPzVpy5LeE576JCwZfWg1ikjUxdK', network='testnet')
         self.assertEqual(b'\xa9\x14\x86\x7f\x84`u\x87\xf7\xc2\x05G@\xc6\xca\xe0\x92\x98\xcc\xbc\xd5(\x87',
                          to.lock_script)
 
-    def test_transaction_output_add_public_key(self):
-        to = Output(1000000000, public_key='0450863AD64A87AE8A2FE83C1AF1A8403CB53F53E486D8511DAD8A04887E5B23522CD470'
-                                           '243453A299FA9E77237716103ABC11A1DF38855ED6F2EE187E9C582BA6')
+    def test_transaction_output_public_key_legacy(self):
+        to = Output(1000000000, public_key='0450863AD64A87AE8A2FE83C1AF1A8403CB53F53E486D8511DAD8A04887E5B23522'
+                                           'CD470243453A299FA9E77237716103ABC11A1DF38855ED6F2EE187E9C582BA6',
+                    witness_type='legacy')
         self.assertEqual(b"v\xa9\x14\x01\tfw`\x06\x95=UgC\x9e^9\xf8j\r';\xee\x88\xac",
                          to.lock_script)
 
-    def test_transaction_output_add_public_key_hash(self):
-        to = Output(1000, public_hash='010966776006953d5567439e5e39f86a0d273bee')
+    def test_transaction_output_public_key(self):
+        to = Output(1000000000, public_key='0450863AD64A87AE8A2FE83C1AF1A8403CB53F53E486D8511DAD8A04887E5B23522'
+                                           'CD470243453A299FA9E77237716103ABC11A1DF38855ED6F2EE187E9C582BA6')
+        self.assertEqual(b"\x00\x14\x01\tfw`\x06\x95=UgC\x9e^9\xf8j\r';\xee",
+                         to.lock_script)
+        self.assertEqual('segwit', to.witness_type)
+
+    def test_transaction_output_public_key_hash(self):
+        to = Output(1000, public_hash='010966776006953d5567439e5e39f86a0d273bee', witness_type='legacy')
         self.assertEqual(b"v\xa9\x14\x01\tfw`\x06\x95=UgC\x9e^9\xf8j\r';\xee\x88\xac",
                          to.lock_script)
 
-    def test_transaction_output_add_script(self):
+    def test_transaction_output_script(self):
         to = Output(1000, lock_script='76a91423e102597c4a99516f851406f935a6e634dbccec88ac')
         self.assertEqual('14GiCdJHj3bznWpcocjcu9ByCmDPEhEoP8', to.address)
 
@@ -140,6 +148,22 @@ def test_transaction_output_value(self):
         self.assertRaisesRegex(ValueError, "Value uses different network \(bitcoin\) then supplied network: testnet",
                                Output, '1 BTC', address=HDKey(network='testnet').address(), network='testnet')
 
+    def test_transaction_output_witness_types(self):
+        k = HDKey(witness_type='segwit')
+        o = Output(10000, k)
+        self.assertEqual(o.witness_type, 'segwit')
+        self.assertEqual(o.script_type, 'p2wpkh')
+
+        k = HDKey(witness_type='legacy')
+        o = Output(10000, k)
+        self.assertEqual(o.witness_type, 'legacy')
+        self.assertEqual(o.script_type, 'p2pkh')
+
+        k = HDKey()
+        o = Output(10000, k)
+        self.assertEqual(o.witness_type, 'segwit')
+        self.assertEqual(o.script_type, 'p2wpkh')
+
 
 class TestTransactions(unittest.TestCase):
     def setUp(self):
@@ -226,7 +250,7 @@ def test_transactions_sign_1(self):
                     keys=pk.public(), network='testnet', witness_type='legacy')
         # key for address mkzpsGwaUU7rYzrDZZVXFne7dXEeo6Zpw2
         pubkey = Key('0391634874ffca219ff5633f814f7f013f7385c66c65c8c7d81e7076a5926f1a75', network='testnet')
-        out = Output(880000, public_hash=pubkey.hash160, network='testnet')
+        out = Output(880000, public_hash=pubkey.hash160, network='testnet', witness_type='legacy')
         t = Transaction([inp], [out], network='testnet')
         t.sign(pk)
         self.assertTrue(t.verify(), msg="Can not verify transaction '%s'")