Improved some methods; fixed some comments

Author: hl75135498756

tvl/crypto/hash.py

@@ -1,6 +1,8 @@
 # Copyright 2023 TropicSquare
 # SPDX-License-Identifier: Apache-2.0
 
+# type: ignore
+
 import crcmod.predefined
 
 _crc16 = crcmod.predefined.mkPredefinedCrcFun("crc-16-buypass")

tvl/crypto/keccak.py

@@ -5,7 +5,7 @@
 from itertools import chain, product
 from math import log2
 from operator import xor
-from typing import List, Protocol, Reversible, Sequence
+from typing import List, Literal, Protocol, Reversible, Sequence
 
 from ..utils import chunked
 from .conversion import bit, bitlist_to_bytes, ints_to_bitlist
@@ -95,14 +95,13 @@ def __call__(self, s: List[List[int]]) -> None:
 
 
 class keccak_p:
-    def __init__(self, b: int, nr: int) -> None:
+    def __init__(self, b: Literal[25, 50, 100, 200, 400, 800, 1600], nr: int) -> None:
         """Keccak-p permutation function as defined in NIST.FIPS.202 3.3
 
         Args:
-            b (int): width of the permutation
+            b (Literal[25, 50, 100, 200, 400, 800, 1600]): width of the permutation
             nr (int): number of rounds
         """
-        assert b in {25, 50, 100, 200, 400, 800, 1600}
         self.b = b
         self.nr = nr
         self._W = self.b // 25
@@ -140,11 +139,11 @@ def __call__(self, s: List[List[int]]) -> None:
             self.round(s, ir)
 
 
-def keccak_f(b: int):
+def keccak_f(b: Literal[25, 50, 100, 200, 400, 800, 1600]):
     """Keccak-f permutation function as defined in NIST.FIPS.202 3.4
 
     Args:
-        b (int): width of the permutation
+        b (Literal[25, 50, 100, 200, 400, 800, 1600]): width of the permutation
 
     Returns:
         the permutation function keccak-f
@@ -289,11 +288,11 @@ def pad10s1(x: int, m: int) -> List[bit]:
     return [bit(1)] + [bit(0)] * j + [bit(1)]
 
 
-def keccak(b: int):
-    """KECCAK sponge as defined in NIST.FIPS.202 6.2
+def keccak(b: Literal[25, 50, 100, 200, 400, 800, 1600]):
+    """KECCAK sponge as defined in NIST.FIPS.202 5.2
 
     Args:
-        b (int): width of the permutation
+        b (Literal[25, 50, 100, 200, 400, 800, 1600]): width of the permutation
     """
 
     def call(r: int):
@@ -303,7 +302,7 @@ def call(r: int):
 
 
 def keccak_c(c: int):
-    """KECCAK[c] sponge as defined in NIST.FIPS.202 6.2
+    """KECCAK[c] sponge as defined in NIST.FIPS.202 5.2
 
     Args:
         c (int): capacity

tvl/messages/message.py

@@ -46,9 +46,8 @@ class _MetaMessage(type):
     """
     Metaclass of the Message class.
 
-    Make sure the field names are unique and move the DataField objects to the
-    class field _FIELD_MODELS so they can used during the instantiation of
-    a new Message object.
+    Make sure the field names are unique
+    and add parameters to the DataField objects
     """
 
     def __new__(
@@ -101,19 +100,16 @@ def specs(cls) -> List[Tuple[str, Type[DataField[Any]], Params]]:
         return sorted(_get_specs(cls), key=lambda x: x[2].priority)
 
     def __init__(self, **kwargs: Any) -> None:
-        for name, type_, params_ in self.specs():
-            if (value := kwargs.get(name)) is None:
-                value = params_.default
-            setattr(self, name, type_(value=value, params=params_))
+        for name, type_, params in self.specs():
+            value = kwargs.get(name, params.default)
+            setattr(self, name, type_(value=value, params=params))
 
     def __eq__(self, __other: Any) -> bool:
         if __other is self:
             return True
-        try:
-            other_to_bytes = __other.to_bytes()
-        except AttributeError:
-            return NotImplemented
-        return other_to_bytes == self.to_bytes()
+        if isinstance(__other, BaseMessage):
+            return __other.to_bytes() == self.to_bytes()
+        return NotImplemented
 
     def __iter__(self) -> Iterator[Tuple[str, DataField[Any]]]:
         yield from self.__dict__.items()