Repository: ethereum/consensus-specs

Issue:
Type hint linting
It should be possibly to lint the pyspec type hints. See [here](https://github.com/ethereum/eth2.0-specs/pull/996#issuecomment-487065832):

> You can run mypy/a linter on the output pyspec file


Discussion:
I've tried to add `mypy` tests - however, the current `SSZType` + `globals()` style type definition is quite incompatible with `mypy`...

I'd say we can solve it with `py-ssz`. :p

We can at least mypy and lint the beacon chain pyspec in CI
Appy `mypy` type hinting check with the current spec: 
```sh
mypy --follow-imports=silent --warn-unused-ignores --ignore-missing-imports --check-untyped-defs --disallow-incomplete-defs --disallow-any-generics -p eth2spec
```
Result:
https://gist.github.com/hwwhww/d000591903690fe9cc4a5eb8a09ee4a9

#### As I can see the problem is that the `def SSZType` + `class SSZObject` tricks are in runtime:
https://github.com/ethereum/eth2.0-specs/blob/f57d6fa28e151570e9c94807e6729a679fc7b1cf/test_libs/pyspec/eth2spec/utils/minimal_ssz.py#L10-L38

For example:
```python
Fork = SSZType({
    # Previous fork version
    'previous_version': 'bytes4',
    # Current fork version
    'current_version': 'bytes4',
    # Fork epoch number
    'epoch': 'uint64',
})
```

And the type checker cannot see functions defined at runtime.


#### On the other hand, `py-ssz` defines a container as a [ssz.sedes.Serializable subclass](https://github.com/ethereum/py-ssz/blob/4430c71d2c9f0a5613c9deb01f65cd3dda545d22/ssz/examples/zoo.py#L23-L29):
```python
class Fork(ssz.Serializable):

    fields = [
        # Previous fork version
        ('previous_version', bytes4),
        # Current fork version
        ('current_version', bytes4),
        # Fork epoch number
        ('epoch', uint64)
    ]

    def __init__(self,
                 previous_version: bytes,
                 current_version: bytes,
                 epoch: Epoch) -> None:
        super().__init__(
            previous_version=previous_version,
            current_version=current_version,
            epoch=epoch,
        )
```
* the fields are knowable to type checkers.
* The `__init__` function rewriting is for type hinting in object initialization.

---

/cc @protolambda @dankrad 
@hwwhww: What is the status of this with the new SSZ typing?
@JustinDrake I was testing on the early version of the new SSZ, I will start to test on the dev branch now.

Resolving patch:
diff --git a/scripts/build_spec.py b/scripts/build_spec.py
--- a/scripts/build_spec.py
+++ b/scripts/build_spec.py
@@ -13,9 +13,10 @@
 
 PHASE0_IMPORTS = '''from typing import (
     Any,
+    Callable,
     Dict,
     List,
-    NewType,
+    Set,
     Tuple,
 )
 
@@ -25,7 +26,8 @@
 )
 from eth2spec.utils.ssz.ssz_typing import (
     # unused: uint8, uint16, uint32, uint128, uint256,
-    uint64, Container, Vector, BytesN
+    uint64, Container, Vector,
+    Bytes4, Bytes32, Bytes48, Bytes96,
 )
 from eth2spec.utils.bls import (
     bls_aggregate_pubkeys,
@@ -38,9 +40,11 @@
 '''
 PHASE1_IMPORTS = '''from typing import (
     Any,
+    Callable,
     Dict,
     List,
-    NewType,
+    Optional,
+    Set,
     Tuple,
 )
 
@@ -52,7 +56,8 @@
 )
 from eth2spec.utils.ssz.ssz_typing import (
     # unused: uint8, uint16, uint32, uint128, uint256,
-    uint64, Container, Vector, BytesN
+    uint64, Container, Vector,
+    Bytes4, Bytes32, Bytes48, Bytes96,
 )
 from eth2spec.utils.bls import (
     bls_aggregate_pubkeys,
@@ -62,13 +67,6 @@
 
 from eth2spec.utils.hash_function import hash
 '''
-NEW_TYPES = {
-    'Slot': 'int',
-    'Epoch': 'int',
-    'Shard': 'int',
-    'ValidatorIndex': 'int',
-    'Gwei': 'int',
-}
 BYTE_TYPES = [4, 32, 48, 96]
 SUNDRY_FUNCTIONS = '''
 def get_ssz_type_by_name(name: str) -> Container:
@@ -77,36 +75,33 @@ def get_ssz_type_by_name(name: str) -> Container:
 
 # Monkey patch validator compute committee code
 _compute_committee = compute_committee
-committee_cache = {}
+committee_cache: Dict[Tuple[Hash, Hash, int, int], List[ValidatorIndex]] = {}
 
 
-def compute_committee(indices: List[ValidatorIndex], seed: Bytes32, index: int, count: int) -> List[ValidatorIndex]:
+def compute_committee(indices: List[ValidatorIndex],  # type: ignore
+                      seed: Hash,
+                      index: int,
+                      count: int) -> List[ValidatorIndex]:
     param_hash = (hash_tree_root(indices), seed, index, count)
 
-    if param_hash in committee_cache:
-        return committee_cache[param_hash]
-    else:
-        ret = _compute_committee(indices, seed, index, count)
-        committee_cache[param_hash] = ret
-        return ret
+    if param_hash not in committee_cache:
+        committee_cache[param_hash] = _compute_committee(indices, seed, index, count)
+    return committee_cache[param_hash]
 
 
 # Monkey patch hash cache
 _hash = hash
-hash_cache = {}
+hash_cache: Dict[bytes, Hash] = {}
 
 
-def hash(x):
-    if x in hash_cache:
-        return hash_cache[x]
-    else:
-        ret = _hash(x)
-        hash_cache[x] = ret
-        return ret
+def hash(x: bytes) -> Hash:
+    if x not in hash_cache:
+        hash_cache[x] = Hash(_hash(x))
+    return hash_cache[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
@@ -120,32 +115,39 @@ def apply_constants_preset(preset: Dict[str, Any]):
 
 
 def objects_to_spec(functions: Dict[str, str],
+                    custom_types: Dict[str, str],
                     constants: Dict[str, str],
                     ssz_objects: Dict[str, str],
                     inserts: Dict[str, str],
                     imports: Dict[str, str],
-                    new_types: Dict[str, str],
-                    byte_types: List[int],
                     ) -> str:
     """
     Given all the objects that constitute a spec, combine them into a single pyfile.
     """
-    new_type_definitions = '\n'.join(['Bytes%s = BytesN[%s]' % (n, n) for n in byte_types])
-    new_type_definitions += '\n' + '\n'.join(['Hash = Bytes32', 'BLSPubkey = Bytes48', 'BLSSignature = Bytes96'])
-    new_type_definitions += \
-        '\n' + '\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: {value}) -> None:\n"
+                f"        ...\n"
+                if value.startswith("uint")
+                else f"class {key}({value}):\n    pass\n"
+                for key, value in custom_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
@@ -171,7 +173,7 @@ def combine_constants(old_constants: Dict[str, str], new_constants: Dict[str, st
     return old_constants
 
 
-def dependency_order_ssz_objects(objects: Dict[str, str]) -> None:
+def dependency_order_ssz_objects(objects: Dict[str, str], custom_types: Dict[str, str]) -> None:
     """
     Determines which SSZ Object is depenedent on which other and orders them appropriately
     """
@@ -180,14 +182,14 @@ def dependency_order_ssz_objects(objects: Dict[str, str]) -> None:
         dependencies = re.findall(r'(: [A-Z][\w[]*)', value)
         dependencies = map(lambda x: re.sub(r'\W|Vector|List|Container|Hash|BLSPubkey|BLSSignature|uint\d+|Bytes\d+|bytes', '', x), dependencies)
         for dep in dependencies:
-            if dep in NEW_TYPES or len(dep) == 0:
+            if dep in custom_types or len(dep) == 0:
                 continue
             key_list = list(objects.keys())
             for item in [dep, key] + key_list[key_list.index(dep)+1:]:
                 objects[item] = objects.pop(item)
 
 
-def combine_ssz_objects(old_objects: Dict[str, str], new_objects: Dict[str, str]) -> Dict[str, str]:
+def combine_ssz_objects(old_objects: Dict[str, str], new_objects: Dict[str, str], custom_types) -> Dict[str, str]:
     """
     Takes in old spec and new spec ssz objects, combines them,
     and returns the newer versions of the objects in dependency order.
@@ -199,7 +201,7 @@ def combine_ssz_objects(old_objects: Dict[str, str], new_objects: Dict[str, str]
             # remove leading variable name
             value = re.sub(r'^class [\w]*\(Container\):\n', '', value)
         old_objects[key] = old_objects.get(key, '') + value
-    dependency_order_ssz_objects(old_objects)
+    dependency_order_ssz_objects(old_objects, custom_types)
     return old_objects
 
 
@@ -211,18 +213,19 @@ def combine_spec_objects(spec0: SpecObject, spec1: SpecObject) -> SpecObject:
     """
     Takes in two spec variants (as tuples of their objects) and combines them using the appropriate combiner function.
     """
-    functions0, constants0, ssz_objects0, inserts0 = spec0
-    functions1, constants1, ssz_objects1, inserts1 = spec1
+    functions0, custom_types0, constants0, ssz_objects0, inserts0 = spec0
+    functions1, custom_types1, constants1, ssz_objects1, inserts1 = spec1
     functions = combine_functions(functions0, functions1)
+    custom_types = combine_constants(custom_types0, custom_types1)
     constants = combine_constants(constants0, constants1)
-    ssz_objects = combine_ssz_objects(ssz_objects0, ssz_objects1)
+    ssz_objects = combine_ssz_objects(ssz_objects0, ssz_objects1, custom_types)
     inserts = combine_inserts(inserts0, inserts1)
-    return functions, constants, ssz_objects, inserts
+    return functions, custom_types, constants, ssz_objects, inserts
 
 
 def build_phase0_spec(sourcefile: str, outfile: str=None) -> Optional[str]:
-    functions, constants, ssz_objects, inserts = get_spec(sourcefile)
-    spec = objects_to_spec(functions, constants, ssz_objects, inserts, PHASE0_IMPORTS, NEW_TYPES, BYTE_TYPES)
+    functions, custom_types, constants, ssz_objects, inserts = get_spec(sourcefile)
+    spec = objects_to_spec(functions, custom_types, constants, ssz_objects, inserts, PHASE0_IMPORTS)
     if outfile is not None:
         with open(outfile, 'w') as out:
             out.write(spec)
@@ -239,7 +242,7 @@ def build_phase1_spec(phase0_sourcefile: str,
     spec_objects = phase0_spec
     for value in [phase1_custody, phase1_shard_data]:
         spec_objects = combine_spec_objects(spec_objects, value)
-    spec = objects_to_spec(*spec_objects, PHASE1_IMPORTS, NEW_TYPES, BYTE_TYPES)
+    spec = objects_to_spec(*spec_objects, PHASE1_IMPORTS)
     if outfile is not None:
         with open(outfile, 'w') as out:
             out.write(spec)
diff --git a/scripts/function_puller.py b/scripts/function_puller.py
--- a/scripts/function_puller.py
+++ b/scripts/function_puller.py
@@ -29,6 +29,7 @@ def get_spec(file_name: str) -> SpecObject:
     inserts = {}
     function_matcher = re.compile(FUNCTION_REGEX)
     inserts_matcher = re.compile(BEGIN_INSERT_REGEX)
+    custom_types = {}
     for linenum, line in enumerate(open(file_name).readlines()):
         line = line.rstrip()
         if pulling_from is None and len(line) > 0 and line[0] == '#' and line[-1] == '`':
@@ -64,7 +65,7 @@ def get_spec(file_name: str) -> SpecObject:
                     ssz_objects[current_name] = ssz_objects.get(current_name, '') + line + '\n'
                 else:
                     functions[current_name] = functions.get(current_name, '') + line + '\n'
-            # Handle constant table entries
+            # Handle constant and custom types table entries
             elif pulling_from is None and len(line) > 0 and line[0] == '|':
                 row = line[1:].split('|')
                 if len(row) >= 2:
@@ -72,12 +73,15 @@ def get_spec(file_name: str) -> SpecObject:
                         row[i] = row[i].strip().strip('`')
                         if '`' in row[i]:
                             row[i] = row[i][:row[i].find('`')]
-                    eligible = True
-                    if row[0][0] not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_':
-                        eligible = False
-                    for c in row[0]:
-                        if c not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_0123456789':
+                    if row[1].startswith('uint') or row[1].startswith('Bytes'):
+                        custom_types[row[0]] = row[1]
+                    else:
+                        eligible = True
+                        if row[0][0] not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_':
                             eligible = False
-                    if eligible:
-                        constants[row[0]] = row[1].replace('**TBD**', '0x1234567890123456789012345678901234567890')
-    return functions, constants, ssz_objects, inserts
+                        for c in row[0]:
+                            if c not in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ_0123456789':
+                                eligible = False
+                        if eligible:
+                            constants[row[0]] = row[1].replace('**TBD**', '0x1234567890123456789012345678901234567890')
+    return functions, custom_types, constants, ssz_objects, inserts


Test patch:
diff --git a/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py b/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
--- a/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
+++ b/test_libs/pyspec/eth2spec/utils/ssz/ssz_typing.py
@@ -47,8 +47,13 @@ def __new__(cls, value, *args, **kwargs):
         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 uint64")
+        return super().__new__(cls, value)
 
 
 class uint128(uint):
@@ -250,7 +255,7 @@ def __init__(self, *args: Iterable):
         # cannot check non-type objects, or parametrized types
         if isinstance(cls.elem_type, type) and not hasattr(cls.elem_type, '__args__'):
             for i, item in enumerate(self.items):
-                if not issubclass(type(item), cls.elem_type):
+                if not issubclass(cls.elem_type, type(item)):
                     raise TypeError("Typed vector cannot hold differently typed value"
                                     " at index %d. Got type: %s, expected type: %s" % (i, type(item), cls.elem_type))
 
@@ -393,11 +398,27 @@ def hash_tree_root(self):
         return hash_tree_root(self, self.__class__)
 
 
+class Bytes4(BytesN):
+    length = 4
+
+
+class Bytes32(BytesN):
+    length = 32
+
+
+class Bytes48(BytesN):
+    length = 48
+
+
+class Bytes96(BytesN):
+    length = 96
+
+
 # SSZ Defaults
 # -----------------------------
 def get_zero_value(typ):
     if is_uint_type(typ):
-        return 0
+        return uint64(0)
     elif is_list_type(typ):
         return []
     elif is_bool_type(typ):
diff --git a/test_libs/pyspec/requirements-testing.txt b/test_libs/pyspec/requirements-testing.txt
--- a/test_libs/pyspec/requirements-testing.txt
+++ b/test_libs/pyspec/requirements-testing.txt
@@ -2,4 +2,5 @@
 pytest>=3.6,<3.7
 ../config_helpers
 flake8==3.7.7
+mypy==0.701
 pytest-cov