From 8b64f37d225bf0d5971735bd8088418dd68b096c Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Mon, 10 Jun 2019 23:16:59 -0400
Subject: [PATCH] Make uint64 be `class` for type hinting

---
 scripts/build_spec.py                         | 36 +++++++++++--------
 .../pyspec/eth2spec/utils/ssz/ssz_typing.py   | 21 ++++++-----
 2 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/scripts/build_spec.py b/scripts/build_spec.py
index 087a88f98..0ede5f1d9 100644
--- a/scripts/build_spec.py
+++ b/scripts/build_spec.py
@@ -15,7 +15,6 @@ PHASE0_IMPORTS = '''from typing import (
     Any,
     Dict,
     List,
-    NewType,
     Tuple,
 )
 
@@ -41,7 +40,6 @@ PHASE1_IMPORTS = '''from typing import (
     Any,
     Dict,
     List,
-    NewType,
     Tuple,
 )
 
@@ -65,11 +63,11 @@ from eth2spec.utils.bls import (
 from eth2spec.utils.hash_function import hash
 '''
 NEW_TYPES = {
-    'Slot': 'int',
-    'Epoch': 'int',
-    'Shard': 'int',
-    'ValidatorIndex': 'int',
-    'Gwei': 'int',
+    'Slot': 'uint64',
+    'Epoch': 'uint64',
+    'Shard': 'uint64',
+    'ValidatorIndex': 'uint64',
+    'Gwei': 'uint64',
 }
 BYTE_TYPES = [4, 32, 48, 96]
 SUNDRY_FUNCTIONS = '''
@@ -79,7 +77,7 @@ def get_ssz_type_by_name(name: str) -> Container:
 
 # Monkey patch validator compute committee code
 _compute_committee = compute_committee
-committee_cache = {}
+committee_cache = {}  # type: Dict[Tuple[Bytes32, Bytes32, ValidatorIndex, int], List[ValidatorIndex]]
 
 
 def compute_committee(indices: List[ValidatorIndex], seed: Bytes32, index: int, count: int) -> List[ValidatorIndex]:
@@ -95,10 +93,10 @@ def compute_committee(indices: List[ValidatorIndex], seed: Bytes32, index: int,
 
 # Monkey patch hash cache
 _hash = hash
-hash_cache = {}
+hash_cache: Dict[bytes, Bytes32] = {}
 
 
-def hash(x):
+def hash(x: bytes) -> Bytes32:
     if x in hash_cache:
         return hash_cache[x]
     else:
@@ -108,7 +106,7 @@ def hash(x):
 
 
 # Access to overwrite spec constants based on configuration
-def apply_constants_preset(preset: Dict[str, Any]):
+def apply_constants_preset(preset: Dict[str, Any]) -> None:
     global_vars = globals()
     for k, v in preset.items():
         global_vars[k] = v
@@ -132,20 +130,28 @@ def objects_to_spec(functions: Dict[str, str],
     """
     Given all the objects that constitute a spec, combine them into a single pyfile.
     """
-    new_type_definitions = \
-        '\n'.join(['''%s = NewType('%s', %s)''' % (key, key, value) for key, value in new_types.items()])
+    new_type_definitions = (
+        '\n\n'.join(
+            [
+                f"class {key}({value}):\n"
+                f"    def __init__(self, _x: uint64) -> None:\n"
+                f"        ...\n"
+                for key, value in new_types.items()
+            ]
+        )
+    )
     functions_spec = '\n\n'.join(functions.values())
     constants_spec = '\n'.join(map(lambda x: '%s = %s' % (x, constants[x]), constants))
     ssz_objects_instantiation_spec = '\n\n'.join(ssz_objects.values())
     ssz_objects_reinitialization_spec = (
-        'def init_SSZ_types():\n    global_vars = globals()\n\n    '
+        'def init_SSZ_types() -> None:\n    global_vars = globals()\n\n    '
         + '\n\n    '.join([re.sub(r'(?!\n\n)\n', r'\n    ', value[:-1]) for value in ssz_objects.values()])
         + '\n\n'
         + '\n'.join(map(lambda x: '    global_vars[\'%s\'] = %s' % (x, x), ssz_objects.keys()))
     )
     spec = (
         imports
-        + '\n' + new_type_definitions
+        + '\n\n' + new_type_definitions
         + '\n\n' + constants_spec
         + '\n\n\n' + ssz_objects_instantiation_spec
         + '\n\n' + functions_spec
diff --git a/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py b/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
index e40c904ca..cde1586c1 100644
--- a/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
+++ b/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
@@ -46,8 +46,13 @@ class uint32(uint):
         return super().__new__(cls, value)
 
 
-# We simply default to uint64. But do give it a name, for readability
-uint64 = NewType('uint64', int)
+class uint64(uint):
+    byte_len = 8
+
+    def __new__(cls, value, *args, **kwargs):
+        if value.bit_length() > 64:
+            raise ValueError("value out of bounds for uint128")
+        return super().__new__(cls, value)
 
 
 class uint128(uint):
@@ -409,12 +414,12 @@ class Bytes96(BytesN):
 # SSZ Defaults
 # -----------------------------
 def get_zero_value(typ):
-    if is_uint_type(typ):
-        return 0
-    elif is_list_type(typ):
+    if is_list_type(typ):
         return []
     elif is_bool_type(typ):
         return False
+    elif is_uint_type(typ):
+        return uint64(0)
     elif is_vector_type(typ):
         return typ()
     elif is_bytesn_type(typ):
@@ -432,12 +437,12 @@ def get_zero_value(typ):
 
 
 def infer_type(obj):
-    if is_uint_type(obj.__class__):
-        return obj.__class__
-    elif isinstance(obj, int):
+    if isinstance(obj, int):
         return uint64
     elif isinstance(obj, list):
         return List[infer_type(obj[0])]
+    elif is_uint_type(obj.__class__):
+        return obj.__class__
     elif isinstance(obj, (Vector, Container, bool, BytesN, bytes)):
         return obj.__class__
     else: