fix spec builder

scripts/phase0/build_spec.py

@@ -16,7 +16,10 @@ from eth2spec.utils.ssz.ssz_impl import (
     hash_tree_root,
     signing_root,
 )
-from eth2spec.utils.ssz.ssz_typing import *
+from eth2spec.utils.ssz.ssz_typing import (
+    uint8, uint16, uint32, uint64, uint128, uint256,
+    Container, Vector, BytesN
+)
 from eth2spec.utils.bls_stub import (
     bls_aggregate_pubkeys,
     bls_verify,
@@ -24,19 +27,18 @@ from eth2spec.utils.bls_stub import (
 )
 from eth2spec.utils.hash_function import hash
 
-
-# stub, will get overwritten by real var
-SLOTS_PER_EPOCH = 64
-
+# Note: 'int' type defaults to being interpreted as a uint64 by SSZ implementation.
 Slot = NewType('Slot', int)  # uint64
 Epoch = NewType('Epoch', int)  # uint64
 Shard = NewType('Shard', int)  # uint64
 ValidatorIndex = NewType('ValidatorIndex', int)  # uint64
 Gwei = NewType('Gwei', int)  # uint64
+
+Bytes4 = BytesN[4]
 Bytes32 = BytesN[32]
-BLSPubkey = NewType('BLSPubkey', BytesN[48])
-BLSSignature = NewType('BLSSignature', BytesN[96])
-Store = None
+Bytes48 = BytesN[48]
+Bytes96 = BytesN[96]
+
 """)
 
     code_lines += function_puller.get_spec(sourcefile)

scripts/phase0/function_puller.py

@@ -1,4 +1,3 @@
-import sys
 from typing import List
 
 
@@ -6,9 +5,11 @@ def get_spec(file_name: str) -> List[str]:
     code_lines = []
     pulling_from = None
     current_name = None
+    # list of current type definition being parsed, or None otherwise
     current_typedef = None
+    # list of (name, definition lines list) tuples.
     type_defs = []
-    for linenum, line in enumerate(open(sys.argv[1]).readlines()):
+    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] == '`':
             current_name = line[line[:-1].rfind('`') + 1: -1]
@@ -19,22 +20,24 @@ def get_spec(file_name: str) -> List[str]:
             if pulling_from is None:
                 pulling_from = linenum
             else:
-                if current_typedef is not None:
-                    assert code_lines[-1] == '}'
-                    code_lines[-1] = ''
-                    code_lines.append('')
                 pulling_from = None
-                current_typedef = None
+                if current_typedef is not None:
+                    type_defs.append((current_name, current_typedef))
+                    current_typedef = None
         else:
-            if pulling_from == linenum and line == '{':
-                code_lines.append('class %s(SSZContainer):' % current_name)
-                current_typedef = current_name
-                type_defs.append(current_name)
-            elif pulling_from is not None:
+            if pulling_from is not None:
                 # Add some whitespace between functions
-                if line[:3] == 'def':
+                if line[:3] == 'def' or line[:5] == 'class':
                     code_lines.append('')
                     code_lines.append('')
+                # Check for SSZ type definitions
+                if len(line) > 18 and line[:6] == 'class ' and line[-12:] == '(Container):':
+                    name = line[6:-12]
+                    # Check consistency with markdown header
+                    assert name == current_name
+                    current_typedef = []
+                if current_typedef is not None:
+                    current_typedef.append(line)
                 code_lines.append(line)
             elif pulling_from is None and len(line) > 0 and line[0] == '|':
                 row = line[1:].split('|')
@@ -54,11 +57,26 @@ def get_spec(file_name: str) -> List[str]:
     # Build type-def re-initialization
     code_lines.append('\n')
     code_lines.append('ssz_types = [\n')
-    for ssz_type_name in type_defs:
-        code_lines.append(f'    {ssz_type_name},\n')
+    for (ssz_type_name, _) in type_defs:
+        code_lines.append(f'    {ssz_type_name},')
     code_lines.append(']')
     code_lines.append('\n')
-    code_lines.append('def get_ssz_type_by_name(name: str) -> SSZType:')
+    code_lines.append('def init_SSZ_types():')
+    code_lines.append('    global_vars = globals()')
+    for ssz_type_name, ssz_type in type_defs:
+        code_lines.append('')
+        for type_line in ssz_type:
+            if len(type_line) > 0:
+                code_lines.append('    ' + type_line)
+    code_lines.append('\n')
+    for (ssz_type_name, _) in type_defs:
+        code_lines.append(f'    global_vars["{ssz_type_name}"] = {ssz_type_name},')
+    code_lines.append('    global_vars["ssz_types"] = [')
+    for (ssz_type_name, _) in type_defs:
+        code_lines.append(f'        {ssz_type_name},')
+    code_lines.append('    ]')
+    code_lines.append('\n')
+    code_lines.append('def get_ssz_type_by_name(name: str) -> Container:')
     code_lines.append('    return globals()[name]')
     code_lines.append('')
     return code_lines