Create ChunkedTransaction class

src/hiero_sdk_python/consensus/topic_message_submit_transaction.py

@@ -1,36 +1,33 @@
 from __future__ import annotations
 
 import math
-from typing import Literal, overload
 
 from hiero_sdk_python.channels import _Channel
-from hiero_sdk_python.client.client import Client
 from hiero_sdk_python.consensus.topic_id import TopicId
 from hiero_sdk_python.crypto.private_key import PrivateKey
 from hiero_sdk_python.executable import _Method
-from hiero_sdk_python.hapi.services import consensus_submit_message_pb2, timestamp_pb2, transaction_pb2
+from hiero_sdk_python.hapi.services import consensus_submit_message_pb2, transaction_pb2
 from hiero_sdk_python.hapi.services.schedulable_transaction_body_pb2 import (
     SchedulableTransactionBody,
 )
+from hiero_sdk_python.transaction.chunked_transaction import ChunkedTransaction
 from hiero_sdk_python.transaction.custom_fee_limit import CustomFeeLimit
-from hiero_sdk_python.transaction.transaction import Transaction
-from hiero_sdk_python.transaction.transaction_id import TransactionId
-from hiero_sdk_python.transaction.transaction_receipt import TransactionReceipt
-from hiero_sdk_python.transaction.transaction_response import TransactionResponse
 
 
-class TopicMessageSubmitTransaction(Transaction):
+class TopicMessageSubmitTransaction(ChunkedTransaction):
     """
     Represents a transaction that submits a message to a Hedera Consensus Service topic.
 
     Allows setting the target topic ID and message, building the transaction body,
     and executing the submission through a network channel.
+
+    Supports automatic chunking for large messages.
     """
 
     def __init__(
         self,
         topic_id: TopicId | None = None,
-        message: str | None = None,
+        message: bytes | str | None = None,
         chunk_size: int | None = None,
         max_chunks: int | None = None,
     ) -> None:
@@ -39,21 +36,34 @@ def __init__(
 
         Args:
             topic_id (TopicId, optional): The ID of the topic.
-            message (str, optional): The message to submit.
-            chunk_size (int, optional): The maximum chunk size in bytes, Default: 1024.
-            max_chunks (int, optional): The maximum number of chunks allowed, Default: 20.
+            message (str or bytes, optional): The message to submit to the topic.
+            chunk_size (int, optional): The maximum chunk size in bytes. Default: 1024.
+            max_chunks (int, optional): The maximum number of chunks allowed. Default: 20.
         """
         super().__init__()
         self.topic_id: TopicId | None = topic_id
-        self.message: str | None = message
-        self.chunk_size: int = chunk_size or 1024
-        self.max_chunks: int = max_chunks or 20
+        self.message: bytes | str | None = message
+        self.chunk_size: int = 1024
+        self.max_chunks: int = 20
+        if chunk_size is not None:
+            self.set_chunk_size(chunk_size)
+
+        if max_chunks is not None:
+            self.set_max_chunks(max_chunks)
 
-        self._current_chunk_index = 0
         self._total_chunks = self.get_required_chunks()
-        self._initial_transaction_id: TransactionId | None = None
-        self._transaction_ids: list[TransactionId] = []
-        self._signing_keys: list[PrivateKey] = []
+
+    def _message_as_bytes(self) -> bytes:
+        """
+        Returns the message encoded as bytes for chunking and protobuf serialization.
+
+        Returns:
+            bytes: The message as bytes.
+        """
+        if self.message is None:
+            return b""
+
+        return self.message.encode("utf-8") if isinstance(self.message, str) else self.message
 
     def get_required_chunks(self) -> int:
         """
@@ -65,7 +75,7 @@ def get_required_chunks(self) -> int:
         if not self.message:
             return 1
 
-        content = self.message.encode("utf-8")
+        content = self._message_as_bytes()
         return math.ceil(len(content) / self.chunk_size)
 
     def set_topic_id(self, topic_id: TopicId) -> TopicMessageSubmitTransaction:
@@ -82,12 +92,12 @@ def set_topic_id(self, topic_id: TopicId) -> TopicMessageSubmitTransaction:
         self.topic_id = topic_id
         return self
 
-    def set_message(self, message: str) -> TopicMessageSubmitTransaction:
+    def set_message(self, message: bytes | str) -> TopicMessageSubmitTransaction:
         """
         Sets the message to submit to the topic.
 
         Args:
-            message (str): The message to submit to the topic.
+            message (str or bytes): The message to submit to the topic.
 
         Returns:
             TopicMessageSubmitTransaction: This transaction instance (for chaining).
@@ -102,16 +112,12 @@ def set_chunk_size(self, chunk_size: int) -> TopicMessageSubmitTransaction:
         Set maximum chunk size in bytes.
 
         Args:
-            chunk_size (int):  The size of each chunk in bytes.
+            chunk_size (int): The size of each chunk in bytes.
 
         Returns:
             TopicMessageSubmitTransaction: This transaction instance (for chaining).
         """
-        self._require_not_frozen()
-        if chunk_size <= 0:
-            raise ValueError("chunk_size must be positive")
-
-        self.chunk_size = chunk_size
+        super().set_chunk_size(chunk_size)
         self._total_chunks = self.get_required_chunks()
         return self
 
@@ -125,11 +131,7 @@ def set_max_chunks(self, max_chunks: int) -> TopicMessageSubmitTransaction:
         Returns:
             TopicMessageSubmitTransaction: This transaction instance (for chaining).
         """
-        self._require_not_frozen()
-        if max_chunks <= 0:
-            raise ValueError("max_chunks must be positive")
-
-        self.max_chunks = max_chunks
+        super().set_max_chunks(max_chunks)
         return self
 
     def set_custom_fee_limits(self, custom_fee_limits: list[CustomFeeLimit]) -> TopicMessageSubmitTransaction:
@@ -160,21 +162,6 @@ def add_custom_fee_limit(self, custom_fee_limit: CustomFeeLimit) -> TopicMessage
         self.custom_fee_limits.append(custom_fee_limit)
         return self
 
-    def _validate_chunking(self) -> None:
-        """
-        Validates that chunk count does not exceed max_chunks.
-
-        Raises:
-            ValueError: If chunk count exceeds `max_chunks`.
-        """
-        required = self.get_required_chunks()
-
-        if self.max_chunks and required > self.max_chunks:
-            raise ValueError(
-                f"Message requires {required} chunks but max_chunks={self.max_chunks}. "
-                f"Increase limit with set_max_chunks()."
-            )
-
     def _build_proto_body(self) -> consensus_submit_message_pb2.ConsensusSubmitMessageTransactionBody:
         """
         Returns the protobuf body for the topic message submit transaction.
@@ -185,10 +172,10 @@ def _build_proto_body(self) -> consensus_submit_message_pb2.ConsensusSubmitMessa
         Raises:
             ValueError: If required fields (message) are missing.
         """
-        if self.message is None or self.message == "":
+        if not self.message:
             raise ValueError("Missing required fields: message.")
 
-        content = self.message.encode("utf-8")
+        content = self._message_as_bytes()
 
         start_index = self._current_chunk_index * self.chunk_size
         end_index = min(start_index + self.chunk_size, len(content))
@@ -215,8 +202,7 @@ def build_transaction_body(self) -> transaction_pb2.TransactionBody:
         Builds and returns the protobuf transaction body for message submission.
 
         Returns:
-            TransactionBody: The protobuf transaction body containing
-                the message submission details.
+            TransactionBody: The protobuf transaction body containing the message submission details.
         """
         consensus_submit_message_body = self._build_proto_body()
         transaction_body = self.build_base_transaction_body()
@@ -247,180 +233,15 @@ def _get_method(self, channel: _Channel) -> _Method:
         """
         return _Method(transaction_func=channel.topic.submitMessage, query_func=None)
 
-    def freeze_with(self, client: Client) -> TopicMessageSubmitTransaction:
-        if self._transaction_body_bytes:
-            return self
-
-        self._resolve_transaction_id(client)
-
-        if not self._transaction_ids:
-            base_timestamp = self.transaction_id.valid_start
-
-            for i in range(self.get_required_chunks()):
-                if i == 0:
-                    if self._initial_transaction_id is None:
-                        self._initial_transaction_id = self.transaction_id
-
-                    chunk_transaction_id = self.transaction_id
-                else:
-                    next_nanos = base_timestamp.nanos + i
-
-                    chunk_valid_start = timestamp_pb2.Timestamp(
-                        seconds=base_timestamp.seconds + next_nanos // 1_000_000_000, nanos=next_nanos % 1_000_000_000
-                    )
-                    chunk_transaction_id = TransactionId(
-                        account_id=self.transaction_id.account_id, valid_start=chunk_valid_start
-                    )
-
-                self._transaction_ids.append(chunk_transaction_id)
-
-        return super().freeze_with(client)
-
-    @overload
-    def execute(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: Literal[True] = True,
-        validate_status: bool = False,
-    ) -> TransactionReceipt: ...
-
-    @overload
-    def execute(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: Literal[False] = False,
-        validate_status: bool = False,
-    ) -> TransactionResponse: ...
-
-    def execute(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: bool = True,
-        validate_status: bool = False,
-    ) -> TransactionReceipt | TransactionResponse:
-        """
-        Executes the topic message submit transaction.
-
-        For multi-chunk transactions, this method will execute all chunks sequentially and return first response.
-
-        Args:
-            client: The client to execute the transaction with.
-            timeout (int | float | None, optional): The total execution timeout (in seconds) for this execution.
-            wait_for_receipt (bool, optional): Whether to wait for consensus and return the receipt.
-                If False, the method returns a TransactionResponse immediately after submission.
-            validate_status: (bool):  Whether the query should automatically validate the transaction status (default False).
-
-        Returns:
-            TransactionReceipt: If wait_for_receipt is True (default)
-            TransactionResponse: If wait_for_receipt is False
-        """
-        # Return the first response as the JS SDK does
-        return self.execute_all(client, timeout, wait_for_receipt, validate_status)[0]
-
-    @overload
-    def execute_all(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: Literal[True] = True,
-        validate_status: bool = False,
-    ) -> list[TransactionReceipt]: ...
-
-    @overload
-    def execute_all(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: Literal[False] = False,
-        validate_status: bool = False,
-    ) -> list[TransactionResponse]: ...
-
-    def execute_all(
-        self,
-        client: Client,
-        timeout: int | float | None = None,
-        wait_for_receipt: bool = True,
-        validate_status: bool = False,
-    ) -> list[TransactionReceipt] | list[TransactionResponse]:
-        """
-        Executes the topic message submit transaction.
-
-        This method will execute all chunks sequentially and return list of all responses.
-
-        Args:
-            client: The client to execute the transaction with.
-            timeout (int | float | None, optional): The total execution timeout (in seconds) for this execution.
-            wait_for_receipt (bool, optional): Whether to wait for consensus and return the receipt.
-                If False, the method returns a TransactionResponse immediately after submission.
-            validate_status: (bool):  Whether the query should automatically validate the transaction status (default False).
-
-        Returns:
-            List[TransactionReceipt]: If wait_for_receipt is True (default)
-            List[TransactionResponse]: If wait_for_receipt is False
-        """
-        self._validate_chunking()
-
-        if self.get_required_chunks() == 1:
-            return [super().execute(client, timeout, wait_for_receipt, validate_status)]
-
-        # Multi-chunk transaction - execute all chunks
-        responses = []
-
-        for chunk_index in range(self.get_required_chunks()):
-            self._current_chunk_index = chunk_index
-
-            if self._transaction_ids and chunk_index < len(self._transaction_ids):
-                self.transaction_id = self._transaction_ids[chunk_index]
-
-            self._transaction_body_bytes.clear()
-            self._signature_map.clear()
-
-            self.freeze_with(client)
-
-            for signing_key in self._signing_keys:
-                super().sign(signing_key)
-
-            # Execute the chunk
-            response = super().execute(client, timeout, wait_for_receipt, validate_status)
-            responses.append(response)
-
-        return responses
-
-    def sign(self, private_key: PrivateKey):
+    def sign(self, private_key: PrivateKey) -> TopicMessageSubmitTransaction:
         """
         Signs the transaction using the provided private key.
 
-        For multi-chunk transactions, this stores the signing key for later use.
-
         Args:
             private_key (PrivateKey): The private key to sign the transaction with.
-        """
-        if private_key not in self._signing_keys:
-            self._signing_keys.append(private_key)
 
+        Returns:
+            TopicMessageSubmitTransaction: This transaction instance (for chaining).
+        """
         super().sign(private_key)
         return self
-
-    @property
-    def body_size_all_chunks(self) -> list[int]:
-        """Returns an array of body sizes for transactions with multiple chunks."""
-        self._require_frozen()
-        sizes = []
-
-        original_index = self._current_chunk_index
-        original_transaction_id = self.transaction_id
-
-        try:
-            for i, transaction_id in enumerate(self._transaction_ids):
-                self._current_chunk_index = i
-                self.transaction_id = transaction_id
-
-                sizes.append(self.body_size)
-        finally:
-            self._current_chunk_index = original_index
-            self.transaction_id = original_transaction_id
-
-        return sizes