feat: Allow PublicKey for TokenCreateTransaction keys

Signed-off-by: Adityarya11 <[email protected]>

Author: Adityarya11

CHANGELOG.md

@@ -24,6 +24,7 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - Added expiration_time, auto_renew_period, auto_renew_account, fee_schedule_key, kyc_key in `TokenCreateTransaction`, `TokenUpdateTransaction` classes
 - Added comprehensive Google-style docstrings to the `CustomFee` class and its methods in `custom_fee.py`.
 - docs: Add `docs/sdk_developers/project_structure.md` to explain repository layout and import paths.
+- feat: Allow `PublicKey` to be used for keys in `TokenCreateTransaction` for non-custodial signing.
 
 ### Changed
 - chore: validate that token airdrop transactions require an available token service on the channel (#632) 

src/hiero_sdk_python/tokens/token_create_transaction.py

@@ -12,7 +12,7 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Optional, Any, List
+from typing import Optional, Any, List, Union
 
 from hiero_sdk_python.Duration import Duration
 from hiero_sdk_python.channels import _Channel
@@ -26,12 +26,15 @@
 from hiero_sdk_python.tokens.token_type import TokenType
 from hiero_sdk_python.tokens.supply_type import SupplyType
 from hiero_sdk_python.account.account_id import AccountId
+from hiero_sdk_python.crypto.public_key import PublicKey
 from hiero_sdk_python.crypto.private_key import PrivateKey
 from hiero_sdk_python.tokens.custom_fee import CustomFee
 
 AUTO_RENEW_PERIOD = Duration(7890000)  # around 90 days in seconds
 DEFAULT_TRANSACTION_FEE = 3_000_000_000
 
+Key = Union[PublicKey, PrivateKey]
+
 @dataclass
 class TokenParams:
     """
@@ -81,14 +84,14 @@ class TokenKeys:
         kyc_key: The KYC key for the token to grant KYC to an account.
     """
 
-    admin_key: Optional[PrivateKey] = None
-    supply_key: Optional[PrivateKey] = None
-    freeze_key: Optional[PrivateKey] = None
-    wipe_key: Optional[PrivateKey] = None
-    metadata_key: Optional[PrivateKey] = None
-    pause_key: Optional[PrivateKey] = None
-    kyc_key: Optional[PrivateKey] = None
-    fee_schedule_key: Optional[PrivateKey] = None
+    admin_key: Optional[Key] = None
+    supply_key: Optional[Key] = None
+    freeze_key: Optional[Key] = None
+    wipe_key: Optional[Key] = None
+    metadata_key: Optional[Key] = None
+    pause_key: Optional[Key] = None
+    kyc_key: Optional[Key] = None
+    fee_schedule_key: Optional[Key] = None
 
 class TokenCreateValidator:
     """Token, key and freeze checks for creating a token as per the proto"""
@@ -368,43 +371,43 @@ def set_memo(self, memo: str) -> "TokenCreateTransaction":
         self._token_params.memo = memo
         return self
 
-    def set_admin_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_admin_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the admin key for the token, which allows updating and deleting the token."""
         self._require_not_frozen()
         self._keys.admin_key = key
         return self
 
-    def set_supply_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_supply_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the supply key for the token, which allows minting and burning tokens."""
         self._require_not_frozen()
         self._keys.supply_key = key
         return self
 
-    def set_freeze_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_freeze_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the freeze key for the token, which allows freezing and unfreezing accounts."""
         self._require_not_frozen()
         self._keys.freeze_key = key
         return self
 
-    def set_wipe_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_wipe_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the wipe key for the token, which allows wiping tokens from an account."""
         self._require_not_frozen()
         self._keys.wipe_key = key
         return self
 
-    def set_metadata_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_metadata_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the metadata key for the token, which allows updating NFT metadata."""
         self._require_not_frozen()
         self._keys.metadata_key = key
         return self
 
-    def set_pause_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_pause_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the pause key for the token, which allows pausing and unpausing the token."""
         self._require_not_frozen()
         self._keys.pause_key = key
         return self
 
-    def set_kyc_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_kyc_key(self, key: Key) -> "TokenCreateTransaction":
         """ Sets the KYC key for the token, which allows granting KYC to an account."""
         self._require_not_frozen()
         self._keys.kyc_key = key
@@ -416,26 +419,35 @@ def set_custom_fees(self, custom_fees: List[CustomFee]) -> "TokenCreateTransacti
         self._token_params.custom_fees = custom_fees
         return self
 
-    def set_fee_schedule_key(self, key: PrivateKey) -> "TokenCreateTransaction":
+    def set_fee_schedule_key(self, key: Key) -> "TokenCreateTransaction":
         """Sets the fee schedule key for the token."""
         self._require_not_frozen()
         self._keys.fee_schedule_key = key
         return self
 
-    def _to_proto_key(self, private_key: Optional[PrivateKey]) -> Optional[basic_types_pb2.Key]:
+    def _to_proto_key(self, key: Optional[Key]) -> Optional[basic_types_pb2.Key]:
         """
-        Helper method to convert a private key to protobuf Key format.
+        Helper method to convert a PrivateKey or PublicKey to protobuf Key format.
 
         Args:
-            private_key (PrivateKey, Optional): The private key to convert, or None
+            key (Key, Optional): The private key or public key to convert, or None
             
         Returns:
-            basic_types_pb2.Key (Optional): The protobuf key or None if private_key is None
+            basic_types_pb2.Key (Optional): The protobuf key or None if key is None
         """
-        if not private_key:
+        if not key:
             return None
 
-        return private_key.public_key()._to_proto()
+        # If PrivateKey, get public key first
+        if isinstance(key, PrivateKey):
+            return key.public_key()._to_proto()
+        
+        # If PublicKey, just convert it...
+        if isinstance(key, PublicKey):
+            return key._to_proto()
+
+        # Handle any other case (though type hinting should prevent this)
+        raise TypeError("Key must be of type PrivateKey or PublicKey")
 
     def freeze_with(self, client) -> "TokenCreateTransaction":
         """

tests/integration/token_create_transaction_e2e_test.py

@@ -3,6 +3,8 @@
 
 from hiero_sdk_python.Duration import Duration
 from hiero_sdk_python.crypto.private_key import PrivateKey
+from hiero_sdk_python.crypto.public_key import PublicKey
+from hiero_sdk_python.transaction.transaction import Transaction
 from hiero_sdk_python.tokens.token_type import TokenType
 from hiero_sdk_python.query.token_info_query import TokenInfoQuery
 from hiero_sdk_python.timestamp import Timestamp
@@ -146,3 +148,73 @@ def test_fungible_token_create_with_fee_schedule_key():
         # TODO (required TokenFeeScheduleUpdateTransaction)
     finally:
         env.close()
+
+@pytest.mark.integration
+def test_token_create_non_custodial_flow():
+    """
+    Tests the full non-custodial flow:
+    1. Operator builds a TX using only a PublicKey.
+    2. Operator gets the transaction bytes.
+    3. User (with the PrivateKey) signs the bytes.
+    4. Operator executes the signed transaction.
+    """
+    
+    env = IntegrationTestEnv()
+    client = env.client
+
+    try:
+        # 1. SETUP: Create a new key pair for the "user"
+        user_private_key = PrivateKey.generate_ed25519()
+        user_public_key = user_private_key.public_key()
+
+        # =================================================================
+        # STEP 1 & 2: OPERATOR (CLIENT) BUILDS THE TRANSACTION
+        # =================================================================
+        
+        tx = (
+            TokenCreateTransaction()
+            .set_token_name("NonCustodialToken")
+            .set_token_symbol("NCT")
+            .set_token_type(TokenType.FUNGIBLE_COMMON)
+            .set_treasury_account_id(client.operator_account_id)
+            .set_initial_supply(100)
+            .set_admin_key(user_public_key)  # <-- The new feature!
+            .set_supply_key(user_public_key) # <-- Using it again
+            .freeze_with(client)
+        )
+
+        tx_bytes = tx.to_bytes()
+
+        # =================================================================
+        # STEP 3: USER (SIGNER) SIGNS THE TRANSACTION
+        # =================================================================
+        
+        tx_from_bytes = Transaction.from_bytes(tx_bytes)
+        tx_from_bytes.sign(user_private_key)
+
+        # =================================================================
+        # STEP 4: OPERATOR (CLIENT) EXECUTES THE SIGNED TX
+        # =================================================================
+        
+        receipt = tx_from_bytes.execute(client)
+        
+        assert receipt is not None
+        token_id = receipt.token_id
+        assert token_id is not None
+        
+        print(f"Successfully created non-custodial token: {token_id}")
+
+        # PROOF: Query the new token and check if the admin key matches
+        token_info = TokenInfoQuery(token_id=token_id).execute(client)
+        
+        assert token_info.admin_key is not None
+        
+        admin_key_bytes = token_info.admin_key.to_bytes_raw()
+        public_key_bytes = user_public_key.to_bytes_raw()
+        
+        assert admin_key_bytes == public_key_bytes
+        print("Integration Test PASSED: Admin key on network matches public key.")
+
+    finally:
+        # Clean up the environment
+        env.close()

tests/unit/test_token_create_transaction.py

@@ -61,6 +61,65 @@ def generate_transaction_id(account_id_proto):
 
 ########### Basic Tests for Building Transactions ###########
 
+def test_set_key_with_public_key(mock_client):
+    """
+    Tests the NEW path: setting a key using a PublicKey.
+    This proves the non-custodial path works.
+    """
+    # 1. Create a real key pair
+    real_private_key = PrivateKey.generate_ed25519()
+    real_public_key = real_private_key.public_key()
+
+    # 2. Build the transaction
+    token_tx = TokenCreateTransaction(
+        TokenParams(
+            token_name="MyToken",
+            token_symbol="MTK",
+            treasury_account_id=AccountId(0, 0, 123),
+            initial_supply=100, # Must be > 0 for fungible
+            token_type=TokenType.FUNGIBLE_COMMON,
+        )
+    )
+    
+    # 3. Use the NEW feature: set a key using a PublicKey
+    token_tx.set_admin_key(real_public_key) 
+
+    # 4. Spy on the _to_proto_key method
+    token_tx._to_proto_key = MagicMock(return_value=real_public_key._to_proto())
+
+    # 5. Freeze (this calls _build_proto_body)
+    token_tx.freeze_with(mock_client)
+
+    # 6. THE PROOF: Assert that _to_proto_key was called 
+    #    with the PublicKey object at least once.
+    token_tx._to_proto_key.assert_any_call(real_public_key)
+    
+def test_set_key_with_invalid_type_raises_error(mock_client):
+    """
+    Tests the SAFETY NET: proves that passing a non-key type fails
+    when the transaction is built.
+    """
+    token_tx = TokenCreateTransaction(
+        TokenParams(
+            token_name="MyToken",
+            token_symbol="MTK",
+            treasury_account_id=AccountId(0, 0, 123),
+            initial_supply=100,
+            token_type=TokenType.FUNGIBLE_COMMON
+        )
+    )
+    
+    # Try to set an invalid key type
+    token_tx.set_admin_key("this is just a string") 
+
+    # Prove that the transaction build fails
+    # with the exact error we want
+    with pytest.raises(TypeError) as e:
+        token_tx.freeze_with(mock_client) # This calls _build_proto_body
+
+    # Check that our specific error was raised
+    assert "Key must be of type PrivateKey or PublicKey" in str(e.value)
+
 # This test uses fixture mock_account_ids as parameter
 def test_build_transaction_body_without_key(mock_account_ids):
     """Test building a token creation transaction body without an admin, supply or freeze key."""
@@ -260,36 +319,36 @@ def test_sign_transaction(mock_account_ids, mock_client):
     private_key.sign.return_value = b"signature"
     private_key.public_key().to_bytes_raw.return_value = b"public_key"
 
-    private_key_admin = MagicMock()
+    private_key_admin = MagicMock(spec=PrivateKey)
     private_key_admin.sign.return_value = b"admin_signature"
     private_key_admin.public_key().to_bytes_raw.return_value = b"admin_public_key"
     private_key_admin.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"admin_public_key")
 
-    private_key_supply = MagicMock()
+    private_key_supply = MagicMock(spec=PrivateKey)
     private_key_supply.sign.return_value = b"supply_signature"
     private_key_supply.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"supply_public_key")
 
-    private_key_freeze = MagicMock()
+    private_key_freeze = MagicMock(spec=PrivateKey)
     private_key_freeze.sign.return_value = b"freeze_signature"
     private_key_freeze.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"freeze_public_key")
 
-    private_key_wipe = MagicMock()
+    private_key_wipe = MagicMock(spec=PrivateKey)
     private_key_wipe.sign.return_value = b"wipe_signature"
     private_key_wipe.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"wipe_public_key")
 
-    private_key_metadata = MagicMock()
+    private_key_metadata = MagicMock(spec=PrivateKey)
     private_key_metadata.sign.return_value = b"metadata_signature"
     private_key_metadata.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"metadata_public_key")
 
-    private_key_pause = MagicMock()
+    private_key_pause = MagicMock(spec=PrivateKey)
     private_key_pause.sign.return_value = b"pause_signature"
     private_key_pause.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"pause_public_key")
 
-    private_key_kyc = MagicMock()
+    private_key_kyc = MagicMock(spec=PrivateKey)
     private_key_kyc.sign.return_value = b"kyc_signature"
     private_key_kyc.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"kyc_public_key")
 
-    private_key_fee_schedule = MagicMock()
+    private_key_fee_schedule = MagicMock(spec=PrivateKey)
     private_key_fee_schedule.sign.return_value = b"fee_schedule_signature"
     private_key_fee_schedule.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"fee_schedule_public_key")
 
@@ -722,36 +781,36 @@ def test_build_and_sign_nft_transaction_to_proto(mock_account_ids, mock_client):
     private_key_private.sign.return_value = b"private_signature"
     private_key_private.public_key().to_bytes_raw.return_value = b"private_public_key"
 
-    private_key_admin = MagicMock()
+    private_key_admin = MagicMock(spec=PrivateKey)
     private_key_admin.sign.return_value = b"admin_signature"
     private_key_admin.public_key().to_bytes_raw.return_value = b"admin_public_key"
     private_key_admin.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"admin_public_key")
 
-    private_key_supply = MagicMock()
+    private_key_supply = MagicMock(spec=PrivateKey)
     private_key_supply.sign.return_value = b"supply_signature"
     private_key_supply.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"supply_public_key")
 
-    private_key_freeze = MagicMock()
+    private_key_freeze = MagicMock(spec=PrivateKey)
     private_key_freeze.sign.return_value = b"freeze_signature"
     private_key_freeze.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"freeze_public_key")
 
-    private_key_wipe = MagicMock()
+    private_key_wipe = MagicMock(spec=PrivateKey)
     private_key_wipe.sign.return_value = b"wipe_signature"
     private_key_wipe.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"wipe_public_key")
 
-    private_key_metadata = MagicMock()
+    private_key_metadata = MagicMock(spec=PrivateKey)
     private_key_metadata.sign.return_value = b"metadata_signature"
     private_key_metadata.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"metadata_public_key")
 
-    private_key_pause = MagicMock()
+    private_key_pause = MagicMock(spec=PrivateKey)
     private_key_pause.sign.return_value = b"pause_signature"
     private_key_pause.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"pause_public_key")
 
-    private_key_kyc = MagicMock()
+    private_key_kyc = MagicMock(spec=PrivateKey)
     private_key_kyc.sign.return_value = b"kyc_signature"
     private_key_kyc.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"kyc_public_key")
 
-    private_key_fee_schedule = MagicMock()
+    private_key_fee_schedule = MagicMock(spec=PrivateKey)
     private_key_fee_schedule.sign.return_value = b"fee_schedule_signature"
     private_key_fee_schedule.public_key()._to_proto.return_value = basic_types_pb2.Key(ed25519=b"fee_schedule_public_key")