diff --git a/ailment/block_walker.py b/ailment/block_walker.py
index 92af66a..411bbf0 100644
--- a/ailment/block_walker.py
+++ b/ailment/block_walker.py
@@ -68,8 +68,8 @@ def walk(self, block: Block) -> None:
             self._handle_stmt(i, stmt, block)
             i += 1
 
-    def walk_statement(self, stmt: Statement):
-        return self._handle_stmt(0, stmt, None)
+    def walk_statement(self, stmt: Statement, block: Block | None = None):
+        return self._handle_stmt(0, stmt, block)
 
     def walk_expression(
         self,