diff --git a/vyper/parser/parser_utils.py b/vyper/parser/parser_utils.py
index 83cdab8b0f5..5998bd3ae9a 100644
--- a/vyper/parser/parser_utils.py
+++ b/vyper/parser/parser_utils.py
@@ -834,7 +834,7 @@ def return_check(self, node: Union[python_ast.AST, List[Any]]) -> bool:
 
 
 def annotate_ast(
-    parsed_ast: python_ast.Module,
+    parsed_ast: Union[python_ast.AST, python.Module],
     source_code: str,
     class_types: Optional[ClassTypes] = None,
 ) -> None: