From 87f6e2753f46f00a2087132e55720a3d654b061e Mon Sep 17 00:00:00 2001
From: antazoey <antyzoa@gmail.com>
Date: Tue, 28 Jan 2025 10:25:55 -0600
Subject: [PATCH] fix: missing block id kwarg

---
 src/ape/managers/chain.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/ape/managers/chain.py b/src/ape/managers/chain.py
index c0d6077f39..b5296a698d 100644
--- a/src/ape/managers/chain.py
+++ b/src/ape/managers/chain.py
@@ -37,8 +37,7 @@
 if TYPE_CHECKING:
     from rich.console import Console as RichConsole
 
-    from ape.types.trace import GasReport, SourceTraceback
-    from ape.types.vm import ContractCode, SnapshotID
+    from ape.types import BlockID, ContractCode, GasReport, SnapshotID, SourceTraceback
 
 
 class BlockContainer(BaseManager):
@@ -716,7 +715,7 @@ class ChainManager(BaseManager):
     _block_container_map: dict[int, BlockContainer] = {}
     _transaction_history_map: dict[int, TransactionHistory] = {}
     _reports: ReportManager = ReportManager()
-    _code: dict[str, dict[str, dict["AddressType", "ContractCode"]]] = {}
+    _code: dict[str, dict[str, dict[AddressType, "ContractCode"]]] = {}
 
     @cached_property
     def contracts(self) -> ContractCache:
@@ -967,12 +966,14 @@ def get_receipt(self, transaction_hash: str) -> ReceiptAPI:
 
         return receipt
 
-    def get_code(self, address: "AddressType") -> "ContractCode":
+    def get_code(
+        self, address: AddressType, block_id: Optional["BlockID"] = None
+    ) -> "ContractCode":
         network = self.provider.network
         if network.is_dev:
             # Avoid caching when dev, as you can manipulate the chain more
             # (and there is isolation).
-            return self.provider.get_code(address)
+            return self.provider.get_code(address, block_id=block_id)
 
         self._code.setdefault(network.ecosystem.name, {})
         self._code[network.ecosystem.name].setdefault(network.name, {})