feat: Allow PublicKey for TokenCreateTransaction keys (hiero-ledger#754)

Signed-off-by: Adityarya11 <arya050411@gmail.com>

Author: Adityarya11

CHANGELOG.md

@@ -8,13 +8,15 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 
 
 ### Added
+
 - Add `examples/topic_id.py` to demonstrate `TopicId` opeartions
 - Add `examples/topic_message.py` to demonstrate `TopicMessage` and `TopicMessageChunk` with local mock data.
 - Added missing validation logic `fee_schedule_key` in integration `token_create_transaction_e2e_test.py` and ``token_update_transaction_e2e_test.py`. 
 - Add `account_balance_query.py` example to demonstrate how to use the CryptoGetAccountBalanceQuery class.
 - Add `examples/token_create_transaction_admin_key.py` demonstrating admin key privileges for token management including token updates, key changes, and deletion (#798)
 - Add `examples/account_info.py` to demonstrate `AccountInfo` opeartions
 - Added `HbarUnit` class and Extend `Hbar` class to handle floating-point numbers
+- feat: Allow `PrivateKey` to be used for keys in `TopicCreateTransaction` for consistency.
 
 
 ### Changed

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
@@ -27,11 +27,14 @@
 from hiero_sdk_python.tokens.supply_type import SupplyType
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.crypto.private_key import PrivateKey
+from hiero_sdk_python.crypto.public_key import PublicKey
 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[PrivateKey, PublicKey]
+
 @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,43 @@ 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 the protobuf Key format.
+
+        This ensures only public keys are serialized:
+        - If a PublicKey is provided, it is used directly.
+        - If a PrivateKey is provided, its corresponding public key is extracted and used.
 
         Args:
-            private_key (PrivateKey, Optional): The private key to convert, or None
+            key (Key, Optional): The PrivateKey or PublicKey to convert.
             
         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.
+            
+        Raises:
+            TypeError: If the provided key is not a PrivateKey, PublicKey, or None.
         """
-        if not private_key:
+        if not key:
             return None
 
-        return private_key.public_key()._to_proto()
+        # If it's a PrivateKey, get its public key first
+        if isinstance(key, PrivateKey):
+            return key.public_key()._to_proto()
+        
+        # If it's already a PublicKey, just convert it
+        if isinstance(key, PublicKey):
+            return key._to_proto()
+
+        # Safety net: This will fail if a non-key is passed
+        raise TypeError("Key must be of type PrivateKey or PublicKey")
+
 
     def freeze_with(self, client) -> "TokenCreateTransaction":
         """

tests/integration/token_create_transaction_e2e_test.py

@@ -6,6 +6,8 @@
 from hiero_sdk_python.response_code import ResponseCode
 from hiero_sdk_python.tokens.custom_fixed_fee import CustomFixedFee
 from hiero_sdk_python.tokens.token_fee_schedule_update_transaction import TokenFeeScheduleUpdateTransaction
+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
@@ -165,3 +167,72 @@ def test_fungible_token_create_with_fee_schedule_key():
 
     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!
+            .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
+        
+        # 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
+        
+        # This is the STRONG assertion:
+        # Compare the bytes of the key from the network
+        # with the bytes of the key we originally used.
+        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
+
+    finally:
+        # Clean up the environment
+        env.close()