Change end line and column for untyped function defs to line of function definition

mypy/fastparse.py

@@ -1014,14 +1014,35 @@ def do_func_def(
                     _dummy_fallback,
                 )
 
-        # End position is always the same.
         end_line = getattr(n, "end_lineno", None)
         end_column = getattr(n, "end_col_offset", None)
 
+        # End line and column span the whole function including its body.
+        # Determine end of the function definition itself.
+        # Fall back to the end of the function definition including its body.
+        def_end_line: int | None = n.lineno
+        def_end_column: int | None = n.col_offset
+
+        returns = n.returns
+        # Use the return type hint if defined.
+        if returns is not None:
+            def_end_line = returns.end_lineno
+            def_end_column = returns.end_col_offset
+        # Otherwise use the last argument in the function definition.
+        elif len(args) > 0:
+            last_arg = args[-1]
+            initializer = last_arg.initializer
+
+            def_end_line = initializer.end_line if initializer else last_arg.end_line
+            def_end_column = initializer.end_column if initializer else last_arg.end_column
+
         self.class_and_function_stack.pop()
         self.class_and_function_stack.append("F")
         body = self.as_required_block(n.body, can_strip=True, is_coroutine=is_coroutine)
         func_def = FuncDef(n.name, args, body, func_type, explicit_type_params)
+        func_def.def_end_line = def_end_line
+        func_def.def_end_column = def_end_column
+
         if isinstance(func_def.type, CallableType):
             # semanal.py does some in-place modifications we want to avoid
             func_def.unanalyzed_type = func_def.type.copy_modified()

mypy/messages.py

@@ -261,6 +261,17 @@ def span_from_context(ctx: Context) -> Iterable[int]:
             assert origin_span is not None
             origin_span = itertools.chain(origin_span, span_from_context(secondary_context))
 
+        end_line = context.end_line if context else -1
+        end_column = context.end_column if context else -1
+
+        # FuncDef's end includes the body, use the def's end information if available.
+        if isinstance(context, FuncDef):
+            end_line = context.def_end_line
+            # column is 1-based, see also format_messages in errors.py
+            end_column = (
+                context.def_end_column + 1 if context.def_end_column is not None else end_column
+            )
+
         self.errors.report(
             context.line if context else -1,
             context.column if context else -1,
@@ -269,8 +280,8 @@ def span_from_context(ctx: Context) -> Iterable[int]:
             file=file,
             offset=offset,
             origin_span=origin_span,
-            end_line=context.end_line if context else -1,
-            end_column=context.end_column if context else -1,
+            end_line=end_line,
+            end_column=end_column,
             code=code,
             allow_dups=allow_dups,
         )

mypy/nodes.py

@@ -781,6 +781,8 @@ class FuncDef(FuncItem, SymbolNode, Statement):
         # Present only when a function is decorated with @typing.datasclass_transform or similar
         "dataclass_transform_spec",
         "docstring",
+        "def_end_line",
+        "def_end_column",
     )
 
     __match_args__ = ("name", "arguments", "type", "body")
@@ -810,6 +812,9 @@ def __init__(
         self.is_mypy_only = False
         self.dataclass_transform_spec: DataclassTransformSpec | None = None
         self.docstring: str | None = None
+        # track the end of the function definition itself
+        self.def_end_line: int | None = None
+        self.def_end_column: int | None = None
 
     @property
     def name(self) -> str:

test-data/unit/check-columns.test

@@ -178,6 +178,79 @@ if int():
     def g(x): # E:5: Function is missing a type annotation
         pass
 
+[case testPrettyFunctionMissingTypeAnnotation]
+# flags: --disallow-untyped-defs --pretty
+def function_with_long_name():
+    pass
+[out]
+main:2:1: error: Function is missing a return type annotation
+    def function_with_long_name():
+    ^
+main:2:1: note: Use "-> None" if function does not return a value
+
+[case testColumnEndFunctionMissingTypeAnnotation]
+# flags: --disallow-untyped-defs --show-error-end
+from typing import Any, Optional
+if int():
+    def f():
+        pass
+
+    def f_no_return(x: int, foo: Optional[str]):
+        pass
+
+    def f_default(x: int, foo: Optional[str] = None):
+        pass
+
+    def f_default_untyped(x, foo = None):
+        pass
+
+    def f_args(x, *args):
+        pass
+
+    def f_kwargs(x, *args, **kwargs):
+        pass
+[builtins fixtures/tuple.pyi]
+[out]
+main:4:5:4:5: error: Function is missing a return type annotation
+main:4:5:4:5: note: Use "-> None" if function does not return a value
+main:7:5:7:47: error: Function is missing a return type annotation
+main:10:5:10:52: error: Function is missing a return type annotation
+main:13:5:13:40: error: Function is missing a type annotation
+main:16:5:16:24: error: Function is missing a type annotation
+main:19:5:19:36: error: Function is missing a type annotation
+
+[case testColumnEndFunctionMissingTypeAnnotationWithReturnType]
+# flags: --disallow-untyped-defs --show-error-end
+def f() -> None:
+    pass
+
+def f_partially_typed(x: int, foo) -> None:
+    pass
+
+def f_untyped(x, foo, *args, **kwargs) -> None:
+    pass
+[builtins fixtures/tuple.pyi]
+[out]
+main:5:1:5:43: error: Function is missing a type annotation for one or more arguments
+main:8:1:8:47: error: Function is missing a type annotation for one or more arguments
+
+[case testColumnEndMultiline]
+# flags: --disallow-untyped-defs --warn-no-return --show-error-end
+def f(
+    x: int,
+    y: int,
+):
+    pass
+
+def g(
+    x: int,
+    y: int,
+) -> int:
+    x = 1
+[out]
+main:2:1:4:11: error: Function is missing a return type annotation
+main:8:1:11:9: error: Missing return statement
+
 [case testColumnNameIsNotDefined]
 ((x)) # E:3: Name "x" is not defined