Support u128, i128, timestamp in TypedData (#1380)

Author: franciszekjob

starknet_py/utils/typed_data.py

@@ -1,10 +1,12 @@
+import re
 from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, List, Optional, Union, cast
 
 from marshmallow import Schema, fields, post_load
 
 from starknet_py.cairo.felt import encode_shortstring
+from starknet_py.constants import FIELD_PRIME
 from starknet_py.hash.hash_method import HashMethod
 from starknet_py.hash.selector import get_selector_from_name
 from starknet_py.net.client_utils import _to_rpc_felt
@@ -84,6 +86,43 @@ class BasicType(Enum):
     CONTRACT_ADDRESS = "ContractAddress"
     CLASS_HASH = "ClassHash"
     BOOL = "bool"
+    U128 = "u128"
+    I128 = "i128"
+    TIMESTAMP = "timestamp"
+
+
+def _encode_value_v1(basic_type: BasicType, value: Union[int, str]) -> Optional[int]:
+    if basic_type in (
+        BasicType.FELT,
+        BasicType.SHORT_STRING,
+        BasicType.CONTRACT_ADDRESS,
+        BasicType.CLASS_HASH,
+    ) and isinstance(value, (int, str)):
+        return parse_felt(value)
+
+    if basic_type in (
+        BasicType.U128,
+        BasicType.TIMESTAMP,
+    ) and isinstance(value, (int, str)):
+        return encode_u128(value)
+
+    if basic_type == BasicType.I128 and isinstance(value, (int, str)):
+        return encode_i128(value)
+
+    return None
+
+
+def _encode_value_v0(
+    basic_type: BasicType,
+    value: Union[int, str],
+) -> Optional[int]:
+    if basic_type in (
+        BasicType.FELT,
+        BasicType.STRING,
+    ) and isinstance(value, (int, str)):
+        return parse_felt(value)
+
+    return None
 
 
 @dataclass(frozen=True)
@@ -128,7 +167,6 @@ def to_dict(self) -> dict:
     def _is_struct(self, type_name: str) -> bool:
         return type_name in self.types
 
-    # pylint: disable=too-many-return-statements
     def _encode_value(
         self,
         type_name: str,
@@ -148,15 +186,18 @@ def _encode_value(
 
         basic_type = BasicType(type_name)
 
-        if (basic_type, self.domain.resolved_revision) in [
-            (BasicType.FELT, Revision.V0),
-            (BasicType.FELT, Revision.V1),
-            (BasicType.STRING, Revision.V0),
-            (BasicType.SHORT_STRING, Revision.V1),
-            (BasicType.CONTRACT_ADDRESS, Revision.V1),
-            (BasicType.CLASS_HASH, Revision.V1),
-        ] and isinstance(value, (int, str)):
-            return parse_felt(value)
+        encoded_value = None
+        if self.domain.resolved_revision == Revision.V0 and isinstance(
+            value, (str, int)
+        ):
+            encoded_value = _encode_value_v0(basic_type, value)
+        elif self.domain.resolved_revision == Revision.V1 and isinstance(
+            value, (str, int)
+        ):
+            encoded_value = _encode_value_v1(basic_type, value)
+
+        if encoded_value is not None:
+            return encoded_value
 
         if basic_type == BasicType.BOOL and isinstance(value, (bool, str, int)):
             return encode_bool(value)
@@ -360,6 +401,54 @@ def encode_bool(value: Union[bool, str, int]) -> int:
     raise ValueError(f"Expected boolean value, got [{value}].")
 
 
+def is_digit_string(s: str, signed=False) -> bool:
+    if signed:
+        return bool(re.fullmatch(r"-?\d+", s))
+    return bool(re.fullmatch(r"\d+", s))
+
+
+def encode_u128(value: Union[str, int]) -> int:
+    def is_in_range(n: int):
+        return 0 <= n < 2**128
+
+    if isinstance(value, str) and value.startswith("0x"):
+        int_value = int(value, 16)
+    elif isinstance(value, str) and is_digit_string(value):
+        int_value = int(value)
+    elif isinstance(value, int):
+        int_value = value
+    else:
+        raise ValueError(f"Value [{value}] is not a valid number.")
+
+    if is_in_range(int_value):
+        return int_value
+    raise ValueError(f"Value [{value}] is out of range for '{BasicType.U128}'.")
+
+
+def encode_i128(value: Union[str, int]) -> int:
+    def is_in_range(n: int):
+        return (n < 2**127) or (n >= (FIELD_PRIME - (2**127)))
+
+    if isinstance(value, str) and value.startswith("0x"):
+        int_value = int(value, 16)
+    elif isinstance(value, str) and is_digit_string(value, True):
+        int_value = int(value)
+    elif isinstance(value, int):
+        int_value = value
+    else:
+        raise ValueError(f"Value [{value}] is not a valid number.")
+
+    if abs(int_value) >= FIELD_PRIME:
+        raise ValueError(
+            f"Values outside the range (-FIELD_PRIME, FIELD_PRIME) are not allowed, [{value}] given."
+        )
+    int_value %= FIELD_PRIME
+
+    if is_in_range(int_value):
+        return int_value
+    raise ValueError(f"Value [{value}] is out of range for '{BasicType.I128}'.")
+
+
 def _get_basic_type_names(revision: Revision) -> List[str]:
     basic_types_v0 = [
         BasicType.FELT,
@@ -373,6 +462,9 @@ def _get_basic_type_names(revision: Revision) -> List[str]:
         BasicType.SHORT_STRING,
         BasicType.CONTRACT_ADDRESS,
         BasicType.CLASS_HASH,
+        BasicType.U128,
+        BasicType.I128,
+        BasicType.TIMESTAMP,
     ]
 
     basic_types = basic_types_v0 if revision == Revision.V0 else basic_types_v1

starknet_py/utils/typed_data_test.py

@@ -16,6 +16,8 @@
     Parameter,
     TypedData,
     encode_bool,
+    encode_i128,
+    encode_u128,
     parse_felt,
 )
 
@@ -320,3 +322,79 @@ def test_encode_bool(value: Union[bool, str, int], expected: int):
 def test_encode_invalid_bool(value: Union[bool, str, int]):
     with pytest.raises(ValueError, match=fr"Expected boolean value, got \[{value}\]."):
         encode_bool(value)
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        (0, 0),
+        (1, 1),
+        (1000000, 1000000),
+        ("0x0", 0),
+        ("0x1", 1),
+        ("0x64", 100),
+        (2 ** 128 - 1, 2 ** 128 - 1),
+    ]
+)
+def test_encode_u128(value: Union[str, int], expected: str):
+    assert encode_u128(value) == expected
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        -1,
+        "-1",
+        1.23,
+        -1.23,
+        "1.23",
+        "-1.23",
+        "example",
+        "0xwrong",
+        2 ** 128,
+        hex(2 ** 128),
+    ]
+)
+def test_encode_invalid_u128(value: Union[str, int]):
+    with pytest.raises(ValueError):
+        encode_u128(value)
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        (0, 0),
+        (1, 1),
+        (1000000, 1000000),
+        ("0x0", 0),
+        ("0x1", 1),
+        ("0x64", 100),
+        (2 ** 127 - 1, 2 ** 127 - 1),
+        (-1, 3618502788666131213697322783095070105623107215331596699973092056135872020480),
+        (-1000000, 3618502788666131213697322783095070105623107215331596699973092056135871020481),
+        (-(2 ** 127), 3618502788666131213697322783095070105452966031871127468241404752419987914753),
+    ]
+)
+def test_encode_i128(value: Union[str, int], expected: str):
+    assert encode_i128(value) == expected
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "example",
+        "0xwrong",
+        1.23,
+        -1.23,
+        "1.23",
+        "-1.23",
+        -2 ** 127 - 1,
+        2 ** 127,
+        str(-2 ** 127 - 1),
+        str(2 ** 127),
+
+    ]
+)
+def test_encode_invalid_i128(value: Union[str, int]):
+    with pytest.raises(ValueError):
+        encode_i128(value)