Handle the dependencies order of dataclass objects

setup.py

@@ -1,3 +1,4 @@
+from enum import Enum, auto
 from setuptools import setup, find_packages, Command
 from setuptools.command.build_py import build_py
 from distutils import dir_util
@@ -14,6 +15,13 @@ class SpecObject(NamedTuple):
     custom_types: Dict[str, str]
     constants: Dict[str, str]
     ssz_objects: Dict[str, str]
+    dataclasses: Dict[str, str]
+
+
+class CodeBlockType(Enum):
+    SSZ = auto()
+    DATACLASS = auto()
+    FUNCTION = auto()
 
 
 def get_spec(file_name: str) -> SpecObject:
@@ -28,8 +36,9 @@ def get_spec(file_name: str) -> SpecObject:
     functions: Dict[str, str] = {}
     constants: Dict[str, str] = {}
     ssz_objects: Dict[str, str] = {}
+    dataclasses: Dict[str, str] = {}
     function_matcher = re.compile(FUNCTION_REGEX)
-    is_ssz = False
+    block_type = CodeBlockType.FUNCTION
     custom_types: Dict[str, str] = {}
     for linenum, line in enumerate(open(file_name).readlines()):
         line = line.rstrip()
@@ -43,20 +52,26 @@ def get_spec(file_name: str) -> SpecObject:
         else:
             # Handle function definitions & ssz_objects
             if pulling_from is not None:
-                # SSZ Object
                 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
-                    is_ssz = True
+                    block_type = CodeBlockType.SSZ
-                # function definition
+                elif line[:10] == '@dataclass':
+                    block_type = CodeBlockType.DATACLASS
                 elif function_matcher.match(line) is not None:
                     current_name = function_matcher.match(line).group(0)
-                    is_ssz = False
+                    block_type = CodeBlockType.FUNCTION
-                if is_ssz:
+
+                if block_type == CodeBlockType.SSZ:
                     ssz_objects[current_name] = ssz_objects.get(current_name, '') + line + '\n'
-                else:
+                elif block_type == CodeBlockType.DATACLASS:
+                    dataclasses[current_name] = dataclasses.get(current_name, '') + line + '\n'
+                elif block_type == CodeBlockType.FUNCTION:
                     functions[current_name] = functions.get(current_name, '') + line + '\n'
+                else:
+                    pass
+
             # Handle constant and custom types table entries
             elif pulling_from is None and len(line) > 0 and line[0] == '|':
                 row = line[1:].split('|')
@@ -75,7 +90,7 @@ def get_spec(file_name: str) -> SpecObject:
                         constants[row[0]] = row[1].replace('**TBD**', '2**32')
                     elif row[1].startswith('uint') or row[1].startswith('Bytes'):
                         custom_types[row[0]] = row[1]
-    return SpecObject(functions, custom_types, constants, ssz_objects)
+    return SpecObject(functions, custom_types, constants, ssz_objects, dataclasses)
 
 
 CONFIG_LOADER = '''
@@ -237,7 +252,7 @@ get_start_shard = cache_this(
     _get_start_shard, lru_size=SLOTS_PER_EPOCH * 3)'''
 
 
-def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
+def objects_to_spec(spec_object: SpecObject, imports: str, fork: str, ordered_class_objects: Dict[str, str]) -> str:
     """
     Given all the objects that constitute a spec, combine them into a single pyfile.
     """
@@ -257,7 +272,7 @@ def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
         if k == "BLS12_381_Q":
             spec_object.constants[k] += "  # noqa: E501"
     constants_spec = '\n'.join(map(lambda x: '%s = %s' % (x, spec_object.constants[x]), spec_object.constants))
-    ssz_objects_instantiation_spec = '\n\n'.join(spec_object.ssz_objects.values())
+    ordered_class_objects_spec = '\n\n'.join(ordered_class_objects.values())
     spec = (
             imports
             + '\n\n' + f"fork = \'{fork}\'\n"
@@ -265,7 +280,7 @@ def objects_to_spec(spec_object: SpecObject, imports: str, fork: str) -> str:
             + '\n' + SUNDRY_CONSTANTS_FUNCTIONS
             + '\n\n' + constants_spec
             + '\n\n' + CONFIG_LOADER
-            + '\n\n' + ssz_objects_instantiation_spec
+            + '\n\n' + ordered_class_objects_spec
             + '\n\n' + functions_spec
             + '\n' + PHASE0_SUNDRY_FUNCTIONS
     )
@@ -291,11 +306,12 @@ ignored_dependencies = [
     'bit', 'boolean', 'Vector', 'List', 'Container', 'BLSPubkey', 'BLSSignature',
     'Bytes1', 'Bytes4', 'Bytes32', 'Bytes48', 'Bytes96', 'Bitlist', 'Bitvector',
     'uint8', 'uint16', 'uint32', 'uint64', 'uint128', 'uint256',
-    'bytes', 'byte', 'ByteList', 'ByteVector'
+    'bytes', 'byte', 'ByteList', 'ByteVector',
+    'Dict', 'dict', 'field',
 ]
 
 
-def dependency_order_ssz_objects(objects: Dict[str, str], custom_types: Dict[str, str]) -> None:
+def dependency_order_class_objects(objects: Dict[str, str], custom_types: Dict[str, str]) -> None:
     """
     Determines which SSZ Object is dependent on which other and orders them appropriately
     """
@@ -332,13 +348,14 @@ 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, custom_types0, constants0, ssz_objects0 = spec0
+    functions0, custom_types0, constants0, ssz_objects0, dataclasses0 = spec0
-    functions1, custom_types1, constants1, ssz_objects1 = spec1
+    functions1, custom_types1, constants1, ssz_objects1, dataclasses1 = 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, custom_types)
-    return SpecObject(functions, custom_types, constants, ssz_objects)
+    dataclasses = combine_functions(dataclasses0, dataclasses1)
+    return SpecObject(functions, custom_types, constants, ssz_objects, dataclasses)
 
 
 fork_imports = {
@@ -354,9 +371,10 @@ def build_spec(fork: str, source_files: List[str]) -> str:
     for value in all_specs[1:]:
         spec_object = combine_spec_objects(spec_object, value)
 
-    dependency_order_ssz_objects(spec_object.ssz_objects, spec_object.custom_types)
+    class_objects = {**spec_object.ssz_objects, **spec_object.dataclasses}
+    dependency_order_class_objects(class_objects, spec_object.custom_types)
 
-    return objects_to_spec(spec_object, fork_imports[fork], fork)
+    return objects_to_spec(spec_object, fork_imports[fork], fork, class_objects)
 
 
 class PySpecCommand(Command):