chore: added type var to chunked transaction

Signed-off-by: Manish Dait <daitmanish88@gmail.com>

Author: manishdait

src/hiero_sdk_python/transaction/chunked_transaction.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Literal, overload
+from typing import Literal, TypeVar, overload
 
 from hiero_sdk_python.client.client import Client
 from hiero_sdk_python.transaction.transaction import Transaction
@@ -10,6 +10,9 @@
 from hiero_sdk_python.transaction.transaction_response import TransactionResponse
 
 
+T = TypeVar("T", bound="ChunkedTransaction")
+
+
 class ChunkedTransaction(Transaction, ABC):
     """
     Abstract base class for transactions that support chunking.
@@ -36,7 +39,7 @@ def __init__(self) -> None:
         self.max_chunks: int = 20
 
     @abstractmethod
-    def _build_proto_body(self):
+    def _build_proto_body(self: T):
         """
         Builds the protobuf body for the current chunk.
 
@@ -52,7 +55,7 @@ def _build_proto_body(self):
         """
         pass
 
-    def set_chunk_size(self, chunk_size: int) -> ChunkedTransaction:
+    def set_chunk_size(self: T, chunk_size: int) -> T:
         """
         Sets the chunk size for this transaction.
 
@@ -73,7 +76,7 @@ def set_chunk_size(self, chunk_size: int) -> ChunkedTransaction:
         self._total_chunks = self.get_required_chunks()
         return self
 
-    def set_max_chunks(self, max_chunks: int) -> ChunkedTransaction:
+    def set_max_chunks(self: T, max_chunks: int) -> T:
         """
         Sets the maximum number of chunks allowed.
 
@@ -93,7 +96,7 @@ def set_max_chunks(self, max_chunks: int) -> ChunkedTransaction:
         self.max_chunks = max_chunks
         return self
 
-    def _validate_chunking(self) -> None:
+    def _validate_chunking(self: T) -> None:
         """
         Validates that the transaction doesn't exceed the maximum number of chunks.
 
@@ -106,7 +109,7 @@ def _validate_chunking(self) -> None:
                 f"Required: {self.get_required_chunks()} Increase limit with set_max_chunks()."
             )
 
-    def freeze_with(self, client: Client) -> ChunkedTransaction:
+    def freeze_with(self: T, client: Client) -> T:
         """
         Freezes the transaction by building transaction bodies for all chunks.
 
@@ -149,14 +152,14 @@ def freeze_with(self, client: Client) -> ChunkedTransaction:
 
             self._transaction_body_bytes[self._transaction_ids[chunk]] = node_bytes
 
-        self._current_chunk_index = 0
+        self._current_chunk_index = None
         self._current_transaction_id_index = 0
 
         return super().freeze_with(client)
 
     @overload
     def execute(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: Literal[True] = True,
@@ -165,15 +168,15 @@ def execute(
 
     @overload
     def execute(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: Literal[False] = False,
         validate_status: bool = False,
     ) -> TransactionResponse: ...
 
     def execute(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: bool = True,
@@ -200,7 +203,7 @@ def execute(
 
     @overload
     def execute_all(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: Literal[True] = True,
@@ -209,15 +212,15 @@ def execute_all(
 
     @overload
     def execute_all(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: Literal[False] = False,
         validate_status: bool = False,
     ) -> list[TransactionResponse]: ...
 
     def execute_all(
-        self,
+        self: T,
         client: Client,
         timeout: int | float | None = None,
         wait_for_receipt: bool = True,
@@ -258,7 +261,7 @@ def execute_all(
         return responses
 
     @property
-    def body_size_all_chunks(self) -> list[int]:
+    def body_size_all_chunks(self: T) -> list[int]:
         """
         Returns an array of body sizes for each chunk in the transaction.
 
@@ -273,6 +276,7 @@ def body_size_all_chunks(self) -> list[int]:
         self._require_frozen()
         sizes = []
 
+        original_chunk_index = self._current_chunk_index
         original_transaction_index = self._current_transaction_id_index
 
         try:
@@ -281,6 +285,7 @@ def body_size_all_chunks(self) -> list[int]:
                 self._current_chunk_index = i
                 sizes.append(self.body_size)
         finally:
+            self._current_chunk_index = original_chunk_index
             self._current_transaction_id_index = original_transaction_index
 
         return sizes

tests/unit/chunked_transaction_test.py

@@ -148,7 +148,8 @@ def test_execute_all_multi_chunk_replays_each_chunk(mock_client, private_key):
 
     assert responses == ["chunk-1", "chunk-2", "chunk-3"]
     assert mock_execute.call_count == 3
-    assert tx._current_chunk_index == 2
+    # ensure reset to none once transaction_body_bytes created
+    assert tx._current_chunk_index is None
 
 
 def test_validate_chunking_allows_required_equal_to_max_chunks():

tests/unit/transaction_test.py

@@ -410,7 +410,7 @@ def test_chunked_tx_return_proper_sizes(file_id, account_id, transaction_id):
     assert large_size > 1024
     # The larger chunked transaction should be bigger than the single-chunk transaction
     assert large_size > small_size
-    assert large_tx._current_chunk_index == 0
+    assert large_tx._current_chunk_index is None
 
 
 def test_chunked_tx_differ_size_if_chunk_are_not_equal(topic_id, account_id, transaction_id):