Properly order types from tables and python code blocks (#4625)

Currently, the pyspec build system forces a specific order where types
defined using table syntax are always emitted before those defined in
Python code blocks, even across forks.

Properly ordering types defined in tables vs code blocks improves spec
clarity as it enables `List` typealiases to be defined in tables that
use an element type that is originally defined in a code block.

Author: etan-status

pysetup/generate_specs.py

@@ -43,6 +43,7 @@
 from pysetup.helpers import (
     combine_spec_objects,
     dependency_order_class_objects,
+    finalized_spec_object,
     objects_to_spec,
     parse_config_vars,
 )
@@ -118,6 +119,7 @@ def build_spec(
     spec_object = all_specs[0]
     for value in all_specs[1:]:
         spec_object = combine_spec_objects(spec_object, value)
+    spec_object = finalized_spec_object(spec_object)
 
     class_objects = {**spec_object.ssz_objects, **spec_object.dataclasses}
 
@@ -127,7 +129,7 @@ def build_spec(
         new_objects = copy.deepcopy(class_objects)
         dependency_order_class_objects(
             class_objects,
-            spec_object.custom_types | spec_object.preset_dep_custom_types,
+            spec_object.custom_types,
         )
 
     return objects_to_spec(preset_name, spec_object, fork, class_objects)

pysetup/helpers.py

@@ -23,8 +23,19 @@ def collect_prev_forks(fork: str) -> list[str]:
 
 
 def requires_mypy_type_ignore(value: str) -> bool:
-    return value.startswith("ByteVector") or (
-        value.startswith("Vector") and any(k in value for k in ["ceillog2", "floorlog2"])
+    return (
+        value.startswith("Bitlist")
+        or value.startswith("ByteVector")
+        or (value.startswith("List") and not re.match(r"^List\[\w+,\s*\w+\]$", value))
+        or (value.startswith("Vector") and any(k in value for k in ["ceillog2", "floorlog2"]))
+    )
+
+
+def gen_new_type_definition(name: str, value: str) -> str:
+    return (
+        f"class {name}({value}):\n    pass"
+        if not requires_mypy_type_ignore(value)
+        else f"class {name}(\n    {value}  # type: ignore\n):\n    pass"
     )
 
 
@@ -41,19 +52,11 @@ def objects_to_spec(
     """
 
     def gen_new_type_definitions(custom_types: dict[str, str]) -> str:
-        return "\n\n".join(
-            [
-                (
-                    f"class {key}({value}):\n    pass\n"
-                    if not requires_mypy_type_ignore(value)
-                    else f"class {key}({value}):  # type: ignore\n    pass\n"
-                )
-                for key, value in custom_types.items()
-            ]
+        return "\n\n\n".join(
+            [gen_new_type_definition(key, value) for key, value in custom_types.items()]
         )
 
     new_type_definitions = gen_new_type_definitions(spec_object.custom_types)
-    preset_dep_new_type_definitions = gen_new_type_definitions(spec_object.preset_dep_custom_types)
 
     # Collect builders with the reversed previous forks
     # e.g. `[bellatrix, altair, phase0]` -> `[phase0, altair, bellatrix]`
@@ -224,10 +227,9 @@ def format_constant(name: str, vardef: VariableDefinition) -> str:
         ssz_dep_constants,
         new_type_definitions,
         constant_vars_spec,
-        # The presets that some SSZ types require. Need to be defined before `preset_dep_new_type_definitions`
+        # The presets that some SSZ types require.
         preset_vars_spec,
         preset_dep_constant_vars_spec,
-        preset_dep_new_type_definitions,
         config_spec,
         # Custom classes which are not required to be SSZ containers.
         classes,
@@ -267,8 +269,6 @@ def combine_dicts(old_dict: dict[str, T], new_dict: dict[str, T]) -> dict[str, T
     "bit",
     "Bitlist",
     "Bitvector",
-    "BLSPubkey",
-    "BLSSignature",
     "boolean",
     "byte",
     "ByteList",
@@ -292,6 +292,8 @@ def combine_dicts(old_dict: dict[str, T], new_dict: dict[str, T]) -> dict[str, T
     "floorlog2",
     "List",
     "Optional",
+    "ProgressiveBitlist",
+    "ProgressiveList",
     "Sequence",
     "Set",
     "Tuple",
@@ -312,10 +314,14 @@ def dependency_order_class_objects(objects: dict[str, str], custom_types: dict[s
     items = list(objects.items())
     for key, value in items:
         dependencies = []
-        for line in value.split("\n"):
-            if not re.match(r"\s+\w+: .+", line):
+        for i, line in enumerate(value.split("\n")):
+            if i == 0:
+                match = re.match(r".+\((.+)\):", line)
+            else:
+                match = re.match(r"\s+\w+: (.+)", line)
+            if not match:
                 continue  # skip whitespace etc.
-            line = line[line.index(":") + 1 :]  # strip of field name
+            line = match.group(1)
             if "#" in line:
                 line = line[: line.index("#")]  # strip of comment
             dependencies.extend(
@@ -349,9 +355,6 @@ def combine_spec_objects(spec0: SpecObject, spec1: SpecObject) -> SpecObject:
     protocols = combine_protocols(spec0.protocols, spec1.protocols)
     functions = combine_dicts(spec0.functions, spec1.functions)
     custom_types = combine_dicts(spec0.custom_types, spec1.custom_types)
-    preset_dep_custom_types = combine_dicts(
-        spec0.preset_dep_custom_types, spec1.preset_dep_custom_types
-    )
     constant_vars = combine_dicts(spec0.constant_vars, spec1.constant_vars)
     preset_dep_constant_vars = combine_dicts(
         spec0.preset_dep_constant_vars, spec1.preset_dep_constant_vars
@@ -366,7 +369,6 @@ def combine_spec_objects(spec0: SpecObject, spec1: SpecObject) -> SpecObject:
         functions=functions,
         protocols=protocols,
         custom_types=custom_types,
-        preset_dep_custom_types=preset_dep_custom_types,
         constant_vars=constant_vars,
         preset_dep_constant_vars=preset_dep_constant_vars,
         preset_vars=preset_vars,
@@ -378,6 +380,41 @@ def combine_spec_objects(spec0: SpecObject, spec1: SpecObject) -> SpecObject:
     )
 
 
+def finalized_spec_object(spec_object: SpecObject) -> SpecObject:
+    all_config_dependencies = {
+        vardef.type_name or vardef.type_hint
+        for vardef in (
+            spec_object.constant_vars
+            | spec_object.preset_dep_constant_vars
+            | spec_object.preset_vars
+            | spec_object.config_vars
+        ).values()
+        if (vardef.type_name or vardef.type_hint) is not None
+    }
+
+    custom_types = {}
+    ssz_objects = spec_object.ssz_objects
+    for name, value in spec_object.custom_types.items():
+        if any(k in name for k in all_config_dependencies):
+            custom_types[name] = value
+        else:
+            ssz_objects[name] = gen_new_type_definition(name, value)
+
+    return SpecObject(
+        functions=spec_object.functions,
+        protocols=spec_object.protocols,
+        custom_types=custom_types,
+        constant_vars=spec_object.constant_vars,
+        preset_dep_constant_vars=spec_object.preset_dep_constant_vars,
+        preset_vars=spec_object.preset_vars,
+        config_vars=spec_object.config_vars,
+        ssz_dep_constants=spec_object.ssz_dep_constants,
+        func_dep_presets=spec_object.func_dep_presets,
+        ssz_objects=ssz_objects,
+        dataclasses=spec_object.dataclasses,
+    )
+
+
 def parse_config_vars(conf: dict[str, str]) -> dict[str, str | list[dict[str, str]]]:
     """
     Parses a dict of basic str/int/list types into a dict for insertion into the spec code.

pysetup/md_to_spec.py

@@ -32,7 +32,6 @@ def __init__(
         self.preset_name = preset_name
 
         self.document_iterator: Iterator[Element] = self._parse_document(file_name)
-        self.all_custom_types: dict[str, str] = {}
         self.current_heading_name: str | None = None
 
         # Use a single dict to hold all SpecObject fields
@@ -44,7 +43,6 @@ def __init__(
             "func_dep_presets": {},
             "functions": {},
             "preset_dep_constant_vars": {},
-            "preset_dep_custom_types": {},
             "preset_vars": {},
             "protocols": {},
             "ssz_dep_constants": {},
@@ -180,8 +178,12 @@ def _process_code_class(self, source: str, cls: ast.ClassDef) -> None:
         if class_name != self.current_heading_name:
             raise Exception(f"class_name {class_name} != current_name {self.current_heading_name}")
 
-        if parent_class:
-            assert parent_class == "Container"
+        if parent_class == "ProgressiveContainer":
+            source = re.sub(
+                r"^(.*ProgressiveContainer.*)$", r"\1  # type: ignore", source, flags=re.MULTILINE
+            )
+        else:
+            assert parent_class is None or parent_class == "Container"
         self.spec["ssz_objects"][class_name] = source
 
     def _process_table(self, table: Table) -> None:
@@ -202,9 +204,21 @@ def _process_table(self, table: Table) -> None:
             if not _is_constant_id(name):
                 # Check for short type declarations
                 if value.startswith(
-                    ("uint", "Bytes", "ByteList", "Union", "Vector", "List", "ByteVector")
+                    (
+                        "uint",
+                        "Bitlist",
+                        "Bitvector",
+                        "ByteList",
+                        "ByteVector",
+                        "Bytes",
+                        "List",
+                        "ProgressiveBitlist",
+                        "ProgressiveList",
+                        "Union",
+                        "Vector",
+                    )
                 ):
-                    self.all_custom_types[name] = value
+                    self.spec["custom_types"][name] = value
                 continue
 
             # It is a constant name and a generalized index
@@ -418,7 +432,6 @@ def _process_html_block(self, html: HTMLBlock) -> None:
 
     def _finalize_types(self) -> None:
         """
-        Processes all_custom_types into custom_types and preset_dep_custom_types.
         Calls helper functions to update KZG and CURDLEPROOFS setups if needed.
         """
         # Update KZG trusted setup if needed
@@ -433,17 +446,6 @@ def _finalize_types(self) -> None:
                 self.spec["constant_vars"], self.preset_name
             )
 
-        # Split all_custom_types into custom_types and preset_dep_custom_types
-        self.spec["custom_types"] = {}
-        self.spec["preset_dep_custom_types"] = {}
-        for name, value in self.all_custom_types.items():
-            if any(k in value for k in self.preset) or any(
-                k in value for k in self.spec["preset_dep_constant_vars"]
-            ):
-                self.spec["preset_dep_custom_types"][name] = value
-            else:
-                self.spec["custom_types"][name] = value
-
     def _build_spec_object(self) -> SpecObject:
         """
         Returns the SpecObject using all collected data.
@@ -456,7 +458,6 @@ def _build_spec_object(self) -> SpecObject:
             func_dep_presets=self.spec["func_dep_presets"],
             functions=self.spec["functions"],
             preset_dep_constant_vars=self.spec["preset_dep_constant_vars"],
-            preset_dep_custom_types=self.spec["preset_dep_custom_types"],
             preset_vars=self.spec["preset_vars"],
             protocols=self.spec["protocols"],
             ssz_dep_constants=self.spec["ssz_dep_constants"],
@@ -496,6 +497,8 @@ def _get_class_info_from_ast(cls: ast.ClassDef) -> tuple[str, str | None]:
         parent_class = base.id
     elif isinstance(base, ast.Subscript):
         parent_class = base.value.id
+    elif isinstance(base, ast.Call):
+        parent_class = base.func.id
     else:
         # NOTE: SSZ definition derives from earlier phase...
         # e.g. `phase0.SignedBeaconBlock`

pysetup/typing.py

@@ -18,7 +18,6 @@ class SpecObject(NamedTuple):
     functions: dict[str, str]
     protocols: dict[str, ProtocolDefinition]
     custom_types: dict[str, str]
-    preset_dep_custom_types: dict[str, str]  # the types that depend on presets
     constant_vars: dict[str, VariableDefinition]
     preset_dep_constant_vars: dict[str, VariableDefinition]
     preset_vars: dict[str, VariableDefinition]

tests/infra/test_md_to_spec.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from pysetup.helpers import finalized_spec_object
 from pysetup.md_to_spec import MarkdownToSpec
 from pysetup.typing import SpecObject
 
@@ -35,7 +36,6 @@ def test_constructor_initializes_fields(dummy_file, dummy_preset, dummy_config):
     assert m2s.config == dummy_config
     assert m2s.preset_name == preset_name
     assert isinstance(m2s.spec, dict)
-    assert isinstance(m2s.all_custom_types, dict)
     assert hasattr(m2s, "document_iterator")
     assert m2s.current_heading_name is None
 
@@ -313,7 +313,7 @@ class PayloadAttributes(object):
     assert "PayloadAttributes" not in spec_obj.dataclasses
 
 
-def test_finalize_types_called_and_updates_custom_types(
+def test_finalized_spec_object_updates_custom_types(
     tmp_path, dummy_preset, dummy_config, monkeypatch
 ):
     # Minimal markdown with a type definition
@@ -324,6 +324,11 @@ def test_finalize_types_called_and_updates_custom_types(
 | ---------------- | -------------- | --------------------------------- |
 | `Slot`           | `uint64`       | a slot number                     |
 | `Epoch`          | `uint64`       | an epoch number                   |
+
+| Name            | Value      |
+| --------------- | ---------- |
+| `GENESIS_SLOT`  | `Slot(0)`  |
+| `GENESIS_EPOCH` | `Epoch(0)` |
 """
     file = tmp_path / "types.md"
     file.write_text(md_content)
@@ -334,18 +339,7 @@ def test_finalize_types_called_and_updates_custom_types(
         preset_name="mainnet",
     )
 
-    # Spy on _finalize_types
-    called = {}
-    orig_finalize_types = m2s._finalize_types
-
-    def spy_finalize_types():
-        called["ran"] = True
-        return orig_finalize_types()
-
-    monkeypatch.setattr(m2s, "_finalize_types", spy_finalize_types)
-
-    spec_obj = m2s.run()
-    assert called.get("ran") is True
-    # After _finalize_types, custom_types should include 'Slot' and 'Epoch'
+    spec_obj = finalized_spec_object(m2s.run())
+    # After finalized_spec_object, custom_types should include 'Slot' and 'Epoch'
     assert spec_obj.custom_types["Slot"] == "uint64"
     assert spec_obj.custom_types["Epoch"] == "uint64"