|
7 | 7 |
|
8 | 8 | from __future__ import annotations |
9 | 9 |
|
10 | | -from hiero_sdk_python.crypto.key import Key as SdkKey |
11 | | -from hiero_sdk_python.crypto.private_key import PrivateKey |
12 | | -from hiero_sdk_python.crypto.public_key import PublicKey |
| 10 | +from hiero_sdk_python.crypto.key import Key |
13 | 11 | from hiero_sdk_python.hapi.services import basic_types_pb2 |
14 | 12 |
|
15 | 13 |
|
16 | | -# Type alias for keys that can be either PrivateKey or PublicKey |
17 | | -Key = PrivateKey | PublicKey | SdkKey |
18 | | - |
19 | | - |
20 | 14 | def key_to_proto(key: Key | None) -> basic_types_pb2.Key | None: |
21 | 15 | """ |
22 | | - Helper function to convert a key (PrivateKey or PublicKey) to protobuf Key format. |
| 16 | + Helper function to convert an SDK key to protobuf Key format. |
23 | 17 |
|
24 | | - This function handles the conversion of SDK key types to protobuf format: |
25 | | - - If a PrivateKey is provided, its corresponding public key is extracted and converted. |
26 | | - - If a PublicKey is provided, it is converted directly to protobuf. |
27 | | - - If None is provided, None is returned. |
| 18 | + This function handles any concrete subclass of Key by delegating to its |
| 19 | + to_proto_key() implementation. If None is provided, None is returned. |
28 | 20 |
|
29 | 21 | Args: |
30 | | - key (Optional[Key]): The key to convert (PrivateKey or PublicKey), or None |
| 22 | + key (Optional[Key]): The key to convert, or None |
31 | 23 |
|
32 | 24 | Returns: |
33 | 25 | basic_types_pb2.Key (Optional): The protobuf key or None if key is None |
34 | 26 |
|
35 | 27 | Raises: |
36 | | - TypeError: If the provided key is not a PrivateKey, PublicKey, or None. |
| 28 | + TypeError: If the provided key is not a Key instance or None. |
37 | 29 | """ |
38 | 30 | if key is None: |
39 | 31 | return None |
40 | 32 |
|
41 | | - # If it's a PrivateKey, get the public key first, then convert to proto |
42 | | - if isinstance(key, PrivateKey): |
43 | | - return key.public_key()._to_proto() |
44 | | - |
45 | | - # If it's a PublicKey, convert directly to proto |
46 | | - if isinstance(key, PublicKey): |
47 | | - return key._to_proto() |
48 | | - |
49 | | - # Any other SDK Key implementation (e.g., KeyList / ThresholdKey wrapper) |
50 | | - if isinstance(key, SdkKey): |
| 33 | + if isinstance(key, Key): |
51 | 34 | return key.to_proto_key() |
52 | 35 |
|
53 | | - raise TypeError("Key must be of type PrivateKey or PublicKey, or another SDK Key implementation") |
| 36 | + raise TypeError("Key must be an instance of hiero_sdk_python.crypto.key.Key or None") |
0 commit comments