diff --git a/verified_cogen/main.py b/verified_cogen/main.py
index d995f99..2246c0a 100644
--- a/verified_cogen/main.py
+++ b/verified_cogen/main.py
@@ -103,7 +103,7 @@ def run_once(
     return len(success_zero_tries), len(success), len(failed), cnt
 
 
-def construct_nagini_rewriter(runner_types: List[str]) -> Optional[Runner]:
+def construct_nagini_rewriter(runner_types: List[str]) -> Optional[Rewriter]:
     runner = None
     for runner_type in runner_types:
         if runner_type == "NaginiRewriter":
@@ -117,7 +117,7 @@ def construct_nagini_rewriter(runner_types: List[str]) -> Optional[Runner]:
     return runner
 
 
-def construct_rewriter(extension: str, runner_types: List[str]) -> Optional[Runner]:
+def construct_rewriter(extension: str, runner_types: List[str]) -> Optional[Rewriter]:
     if extension == "py":
         return construct_nagini_rewriter(runner_types)
     if runner_types: