diff --git a/gooey/python_bindings/source_parser.py b/gooey/python_bindings/source_parser.py
index a979089..d3174d6 100644
--- a/gooey/python_bindings/source_parser.py
+++ b/gooey/python_bindings/source_parser.py
@@ -135,10 +135,13 @@ def extract_parser(modulepath):
   main_func = get_nodes_by_containing_attr(funcs, 'main')[0]
   parse_args_assignment = get_nodes_by_containing_attr(main_func.body, 'parse_args')[0]
 
-  # ast seemingly incorrectly reports the line no of ast.If objects (bug?)
-  # thus, if we end on an IF statement, decriment the line number by one.
-  final_line = main_func.body[-1]
-  restart_line = final_line.lineno if not isinstance(final_line, ast.If) else final_line.lineno - 1
+  # ast reports the line no of a block structure as the start of the structure,
+  # not the end, so we look for the line no of the next node after main()
+  # and use that as the end of the main() function.
+  try:
+    restart_line = nodes.body[nodes.body.index(main_func)+1].lineno - 1
+  except IndexError:
+    restart_line = len(source)
 
   module_source = format_source_to_return_parser(
     source,