From 5a1018ef37e5d468efac466a321a0963ba07b06d Mon Sep 17 00:00:00 2001 From: luffy-orf Date: Wed, 17 Jun 2026 22:40:58 +0530 Subject: [PATCH 1/3] fix(checker): narrow missing return type error span to function header --- .../checker/func_checker.py | 4 +-- .../src/guppylang_internals/span.py | 30 ++++++++++++++++++- .../misc_errors/return_not_annotated.err | 4 +-- .../return_not_annotated_none1.err | 2 -- .../return_not_annotated_none2.err | 5 +--- 5 files changed, 33 insertions(+), 12 deletions(-) diff --git a/guppylang-internals/src/guppylang_internals/checker/func_checker.py b/guppylang-internals/src/guppylang_internals/checker/func_checker.py index c5904480f..85e4608f8 100644 --- a/guppylang-internals/src/guppylang_internals/checker/func_checker.py +++ b/guppylang-internals/src/guppylang_internals/checker/func_checker.py @@ -25,6 +25,7 @@ from guppylang_internals.error import GuppyError from guppylang_internals.experimental import check_capturing_closures_enabled from guppylang_internals.nodes import CheckedNestedFunctionDef, NestedFunctionDef +from guppylang_internals.span import function_header_span from guppylang_internals.tys.param import Parameter, TypeParam from guppylang_internals.tys.parsing import ( TypeParsingCtx, @@ -306,8 +307,7 @@ def check_signature( UnsupportedError(func_def.args.defaults[0], "Default arguments") ) if func_def.returns is None: - err = MissingReturnAnnotationError(func_def) - # TODO: Error location is incorrect + err = MissingReturnAnnotationError(function_header_span(func_def)) if all(r.value is None for r in return_nodes_in_ast(func_def)): err.add_sub_diagnostic( MissingReturnAnnotationError.ReturnNone(None, func_def.name) diff --git a/guppylang-internals/src/guppylang_internals/span.py b/guppylang-internals/src/guppylang_internals/span.py index b35366d9d..2b3bbf93d 100644 --- a/guppylang-internals/src/guppylang_internals/span.py +++ b/guppylang-internals/src/guppylang_internals/span.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from typing import TypeAlias -from guppylang_internals.ast_util import get_file, get_line_offset +from guppylang_internals.ast_util import get_file, get_line_offset, get_source from guppylang_internals.error import InternalGuppyError from guppylang_internals.ipython_inspect import normalize_ipython_dummy_files @@ -126,6 +126,34 @@ def to_span(x: ToSpan) -> Span: return Span(start, end) +def function_header_span(func_def: ast.FunctionDef) -> Span: + """Returns a span covering only the function header up to and including `:`.""" + start = to_span(func_def).start + source = get_source(func_def) + file = get_file(func_def) + line_offset = get_line_offset(func_def) + if source is None or file is None or line_offset is None: + if func_def.body: + return Span(start, to_span(func_def.body[0]).start) + return Span(start, to_span(func_def).end) + + lines = source.splitlines() + line_idx = func_def.lineno - 1 + paren_depth = 0 + for i in range(line_idx, len(lines)): + col_begin = func_def.col_offset if i == line_idx else 0 + for col in range(col_begin, len(lines[i])): + char = lines[i][col] + if char == "(": + paren_depth += 1 + elif char == ")": + paren_depth -= 1 + elif char == ":" and paren_depth == 0: + return Span(start, Loc(file, i + line_offset, col + 1)) + + raise InternalGuppyError("function_header_span: Could not find header colon") + + #: List of source lines in a file SourceLines: TypeAlias = list[str] diff --git a/tests/error/misc_errors/return_not_annotated.err b/tests/error/misc_errors/return_not_annotated.err index 9e972f9ec..4dc7db99e 100644 --- a/tests/error/misc_errors/return_not_annotated.err +++ b/tests/error/misc_errors/return_not_annotated.err @@ -3,8 +3,6 @@ Error: Missing type annotation (at $FILE:5:0) 3 | 4 | @compile_guppy 5 | def foo(x: bool): - | ^^^^^^^^^^^^^^^^^ -6 | return x - | ^^^^^^^^^^^^ Return type must be annotated + | ^^^^^^^^^^^^^^^^^ Return type must be annotated Guppy compilation failed due to 1 previous error diff --git a/tests/error/misc_errors/return_not_annotated_none1.err b/tests/error/misc_errors/return_not_annotated_none1.err index 0c4e5563f..f87570522 100644 --- a/tests/error/misc_errors/return_not_annotated_none1.err +++ b/tests/error/misc_errors/return_not_annotated_none1.err @@ -3,8 +3,6 @@ Error: Missing type annotation (at $FILE:5:0) 3 | 4 | @compile_guppy 5 | def foo(): - | ^^^^^^^^^^ -6 | return | ^^^^^^^^^^ Return type must be annotated Help: Looks like `foo` doesn't return anything. Consider annotating it with `-> diff --git a/tests/error/misc_errors/return_not_annotated_none2.err b/tests/error/misc_errors/return_not_annotated_none2.err index b6974fc34..f87570522 100644 --- a/tests/error/misc_errors/return_not_annotated_none2.err +++ b/tests/error/misc_errors/return_not_annotated_none2.err @@ -3,10 +3,7 @@ Error: Missing type annotation (at $FILE:5:0) 3 | 4 | @compile_guppy 5 | def foo(): - | ^^^^^^^^^^ - | ... -7 | return x - | ^^^^^^^^^^^^^^^^ Return type must be annotated + | ^^^^^^^^^^ Return type must be annotated Help: Looks like `foo` doesn't return anything. Consider annotating it with `-> None`. From 6fffa7caed0e59d7cb071b118f037eb294d0aaec Mon Sep 17 00:00:00 2001 From: luffy-orf Date: Thu, 18 Jun 2026 21:35:55 +0530 Subject: [PATCH 2/3] fix(checker): address review feedback on function header span --- .../src/guppylang_internals/span.py | 14 +++--- tests/test_function_header_span.py | 43 +++++++++++++++++++ 2 files changed, 50 insertions(+), 7 deletions(-) create mode 100644 tests/test_function_header_span.py diff --git a/guppylang-internals/src/guppylang_internals/span.py b/guppylang-internals/src/guppylang_internals/span.py index 2b3bbf93d..6cd6a351b 100644 --- a/guppylang-internals/src/guppylang_internals/span.py +++ b/guppylang-internals/src/guppylang_internals/span.py @@ -132,18 +132,18 @@ def function_header_span(func_def: ast.FunctionDef) -> Span: source = get_source(func_def) file = get_file(func_def) line_offset = get_line_offset(func_def) - if source is None or file is None or line_offset is None: - if func_def.body: - return Span(start, to_span(func_def.body[0]).start) - return Span(start, to_span(func_def).end) + # `check_signature` is only called on AST nodes that have been processed by + # `annotate_location`, so source metadata is always available. + assert source is not None + assert file is not None + assert line_offset is not None lines = source.splitlines() line_idx = func_def.lineno - 1 paren_depth = 0 - for i in range(line_idx, len(lines)): + for i, line in enumerate(lines[line_idx:], start=line_idx): col_begin = func_def.col_offset if i == line_idx else 0 - for col in range(col_begin, len(lines[i])): - char = lines[i][col] + for col, char in enumerate(line[col_begin:], start=col_begin): if char == "(": paren_depth += 1 elif char == ")": diff --git a/tests/test_function_header_span.py b/tests/test_function_header_span.py new file mode 100644 index 000000000..c9d6ca384 --- /dev/null +++ b/tests/test_function_header_span.py @@ -0,0 +1,43 @@ +import ast + +import pytest +from guppylang_internals.ast_util import annotate_location +from guppylang_internals.span import function_header_span, to_span + + +def _parse_func(source: str) -> ast.FunctionDef: + node = ast.parse(source).body[0] + assert isinstance(node, ast.FunctionDef) + annotate_location(node, source, "test.py", 1) + return node + + +def _header_text(func_def: ast.FunctionDef) -> str: + span = function_header_span(func_def) + source = func_def.source # type: ignore[attr-defined] + lines = source.splitlines() + if span.is_multiline: + parts = [lines[span.start.line - 1][span.start.column :]] + parts.extend( + lines[line_no - 1] + for line_no in range(span.start.line + 1, span.end.line) + ) + parts.append(lines[span.end.line - 1][: span.end.column]) + return "\n".join(parts) + line = lines[span.start.line - 1] + return line[span.start.column : span.end.column] + + +@pytest.mark.parametrize( + ("source", "expected_header"), + [ + ("def foo():\n return", "def foo():"), + ("def foo(x: bool):\n return x", "def foo(x: bool):"), + ("def foo(\n x: bool,\n):\n return x", "def foo(\n x: bool,\n):"), + ("def foo() :\n return", "def foo() :"), + ], +) +def test_function_header_span(source: str, expected_header: str) -> None: + func_def = _parse_func(source) + assert _header_text(func_def) == expected_header + assert function_header_span(func_def).end <= to_span(func_def).end From 86468e291576481d4ead9ba2064a4c0a968312d4 Mon Sep 17 00:00:00 2001 From: luffy-orf Date: Thu, 18 Jun 2026 21:38:27 +0530 Subject: [PATCH 3/3] style: format function header span test --- tests/test_function_header_span.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_function_header_span.py b/tests/test_function_header_span.py index c9d6ca384..c27b18254 100644 --- a/tests/test_function_header_span.py +++ b/tests/test_function_header_span.py @@ -19,8 +19,7 @@ def _header_text(func_def: ast.FunctionDef) -> str: if span.is_multiline: parts = [lines[span.start.line - 1][span.start.column :]] parts.extend( - lines[line_no - 1] - for line_no in range(span.start.line + 1, span.end.line) + lines[line_no - 1] for line_no in range(span.start.line + 1, span.end.line) ) parts.append(lines[span.end.line - 1][: span.end.column]) return "\n".join(parts)