review feedback adjustments

Author: felix314159

src/ethereum_test_forks/__init__.py

@@ -1,7 +1,5 @@
 """Ethereum test fork definitions."""
 
-from typing import Literal
-
 from .base_fork import Fork, ForkAttribute
 from .forks.forks import (
     ArrowGlacier,

src/ethereum_test_types/blob_types.py

@@ -20,20 +20,6 @@
     Prague,
 )
 
-
-def fork_string_to_object(fork_name: str) -> type[Cancun] | type[Prague] | type[Osaka]:
-    """Take a fork string and return the respective fork as object."""
-    fork_name = fork_name.lower()
-
-    if fork_name == "cancun":
-        return Cancun
-    if fork_name == "prague":
-        return Prague
-    if fork_name == "osaka":
-        return Osaka
-    raise ValueError(f"Fork {fork_name} has not yet been implemented in this function.")
-
-
 CACHED_BLOBS_DIRECTORY: Path = (
     Path(platformdirs.user_cache_dir("ethereum-execution-spec-tests")) / "cached_blobs"
 )
@@ -49,7 +35,7 @@ class Blob(CamelModel):
 
     versioned_hash: Hash
     name: str  # blob_<fork>_<seed>
-    fork: str
+    fork: Fork
     seed: int
     timestamp: int
 
@@ -59,7 +45,7 @@ class Blob(CamelModel):
     def trusted_setup(cls):
         """Set trusted setup if it is not already set."""
         if cls._trusted_setup is None:
-            trusted_setup_path = Path(realpath(__file__)).parent / "blob_trusted_setup.txt"
+            trusted_setup_path = Path(realpath(__file__)).parent / "kzg_trusted_setup.txt"
             trusted_setup = ckzg.load_trusted_setup(str(trusted_setup_path), 0)
             print(f"{trusted_setup_path} has been loaded")
             cls._trusted_setup = trusted_setup
@@ -159,7 +145,7 @@ def get_proof(fork: Fork, data: Bytes) -> List[Bytes] | Bytes:
                 return proofs
 
             raise AssertionError(
-                f"get_proof() has not been implemented yet for fork: {fork_str}. "
+                f"get_proof() has not been implemented yet for fork: {fork.name()}."
                 f"Got amount of cell proofs {amount_cell_proofs} but expected 128."
             )
 
@@ -179,14 +165,12 @@ def get_cells(fork: Fork, data: Bytes) -> List[Bytes] | None:
                 return cells  # List[bytes]
 
             raise AssertionError(
-                f"get_cells() has not been implemented yet for fork: {fork_str}. Got amount of "
+                f"get_cells() has not been implemented yet for fork: {fork.name()}. Got amount of "
                 f"cell proofs {amount_cell_proofs} but expected 128."
             )
 
         assert fork.supports_blobs(), f"Provided fork {fork.name()} does not support blobs!"
 
-        fork_str: str = fork.name().lower()
-
         # if this blob already exists then load from file
         blob_location: Path = Blob.get_filepath(fork, seed)
         if blob_location.exists():
@@ -211,7 +195,7 @@ def get_cells(fork: Fork, data: Bytes) -> List[Bytes] | None:
             cells=cells,
             versioned_hash=versioned_hash,
             name=name,
-            fork=fork_str,
+            fork=fork,
             seed=seed,
             timestamp=timestamp,
         )
@@ -253,7 +237,7 @@ def from_file(file_name: str) -> "Blob":
     def write_to_file(self):
         """Take a blob object, serialize it and write it to disk as json."""
         json_str = self.model_dump_json()
-        output_location = Blob.get_filepath(fork_string_to_object(self.fork), self.seed)
+        output_location = Blob.get_filepath(self.fork, self.seed)
 
         # warn if existing static_blob gets overwritten
         if output_location.exists():
@@ -264,11 +248,10 @@ def write_to_file(self):
 
     def verify_cell_kzg_proof_batch(self, cell_indices: list) -> bool:
         """Check whether all cell proofs are valid and returns True only if that is the case."""
-        fork_obj = fork_string_to_object(self.fork)
-        amount_cell_proofs: int = cast(int, fork_obj.get_blob_constant("AMOUNT_CELL_PROOFS"))
+        amount_cell_proofs: int = cast(int, self.fork.get_blob_constant("AMOUNT_CELL_PROOFS"))
 
         assert amount_cell_proofs > 0, (
-            f"verify_cell_kzg_proof_batch() is not available for your fork: {self.fork}."
+            f"verify_cell_kzg_proof_batch() is not available for your fork: {self.fork.name()}."
         )
 
         assert self.cells is not None, "self.cells is None, critical error."
@@ -298,11 +281,10 @@ def delete_cells_then_recover_them(self, deletion_indices: list[int]):
         the ckzg recovery mechanism is used to repair the missing cells.
         If no assertion is triggered the reconstruction was successful.
         """
-        fork_obj = fork_string_to_object(self.fork)
-        amount_cell_proofs: int = cast(int, fork_obj.get_blob_constant("AMOUNT_CELL_PROOFS"))
+        amount_cell_proofs: int = cast(int, self.fork.get_blob_constant("AMOUNT_CELL_PROOFS"))
 
         assert amount_cell_proofs > 0, (
-            f"delete_cells_then_recover_them() is not available for fork: {self.fork}"
+            f"delete_cells_then_recover_them() is not available for fork: {self.fork.name()}"
         )
 
         assert self.cells is not None, "self.cells is None, critical problem."
@@ -377,8 +359,7 @@ def corrupt_byte(b: bytes) -> Bytes:
             return Bytes(bytes([b[0] ^ 0xFF]))
 
         # osaka and later
-        fork_obj = fork_string_to_object(self.fork)
-        amount_cell_proofs: int = cast(int, fork_obj.get_blob_constant("AMOUNT_CELL_PROOFS"))
+        amount_cell_proofs: int = cast(int, self.fork.get_blob_constant("AMOUNT_CELL_PROOFS"))
         if amount_cell_proofs > 0:
             assert isinstance(self.proof, list), (
                 "proof was expected to be a list but it isn't"
@@ -406,7 +387,7 @@ def corrupt_byte(b: bytes) -> Bytes:
 
         # pre-osaka (cancun and prague)
         assert amount_cell_proofs == 0, (
-            f"You need to adjust corrupt_proof to handle fork {self.fork}"
+            f"You need to adjust corrupt_proof to handle fork {self.fork.name()}"
         )
         assert isinstance(self.proof, Bytes), "proof was expected to be Bytes but it isn't"
 
@@ -426,8 +407,6 @@ def corrupt_byte(b: bytes) -> Bytes:
 # mytestfork.fork_at(timestamp=15000)  # returns prague
 # TODO: call this once only at blob creation (not when reading blob from file) + remove timestamp from json
 
-# TODO: in test_blob_txs_full.py cancun test: how to convert Hash to bytes?
-
 # TODO: test all functions again
 
 # --------- generate static blobs ------------
@@ -456,7 +435,9 @@ def corrupt_byte(b: bytes) -> Bytes:
 # duration_ms = (end - start) * 1000
 # print(f"Generated and wrote {amount_of_blobs} blobs to disk in: {duration_ms:.3f} ms")
 
+
 myosaka: Fork = Osaka
+print(f"Fork test print: {myosaka.name()}")
 myprague: Fork = Prague
 mycancun: Fork = Cancun
 myseed: int = 1337  # fork+seed is the unique ID of a blob