From 6dee499a5d849f96ae6585247a859302e7ff4f79 Mon Sep 17 00:00:00 2001
From: Matthias Schoettle <git@mattsch.com>
Date: Mon, 6 May 2024 13:43:09 -0400
Subject: [PATCH] Determine end line and column of function itself

---
 mypy/fastparse.py                 | 21 +++++++++++++++++++++
 mypy/messages.py                  |  7 +++----
 test-data/unit/check-columns.test | 13 +++++++++----
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/mypy/fastparse.py b/mypy/fastparse.py
index e208e4d0b7d9..2e0d976a54e1 100644
--- a/mypy/fastparse.py
+++ b/mypy/fastparse.py
@@ -983,10 +983,31 @@ def do_func_def(
         end_line = getattr(n, "end_lineno", None)
         end_column = getattr(n, "end_col_offset", None)
 
+        # 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)
+        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()
diff --git a/mypy/messages.py b/mypy/messages.py
index 424ae1546e9e..64c312612365 100644
--- a/mypy/messages.py
+++ b/mypy/messages.py
@@ -256,7 +256,6 @@ 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))
 
-        column = context.column if context else -1
         end_line = context.end_line if context else -1
         end_column = context.end_column if context else -1
 
@@ -264,13 +263,13 @@ def span_from_context(ctx: Context) -> Iterable[int]:
         # this avoids errors being reported in IDEs for the whole function
         # TODO: figure out if it's possible to find the end of the function definition line
         if isinstance(context, FuncDef):
-            end_line = context.line
+            end_line = context.def_end_line
             # column is 1-based, see also format_messages in errors.py
-            end_column = column + 1
+            end_column = context.def_end_column + 1 if context.def_end_column else end_column
 
         self.errors.report(
             context.line if context else -1,
-            column,
+            context.column if context else -1,
             msg,
             severity=severity,
             file=file,
diff --git a/test-data/unit/check-columns.test b/test-data/unit/check-columns.test
index 0a376426fafb..600cabcce9b8 100644
--- a/test-data/unit/check-columns.test
+++ b/test-data/unit/check-columns.test
@@ -180,15 +180,20 @@ if int():
 
 [case testColumnEndFunctionMissingTypeAnnotation]
 # flags: --disallow-untyped-defs --show-error-end
+from typing import Any, Optional
 if int():
-    def f(x: int):
+    def f(x: int, foo: Optional[str]):
         pass
 
-    def g(x):
+    def g(x: int, foo: Optional[str] = None):
+        pass
+
+    def h(x, foo = None):
         pass
 [out]
-main:3:5:3:5: error: Function is missing a return type annotation
-main:6:5:6:5: error: Function is missing a type annotation
+main:4:5:4:37: error: Function is missing a return type annotation
+main:7:5:7:44: error: Function is missing a return type annotation
+main:10:5:10:24: error: Function is missing a type annotation
 
 [case testColumnNameIsNotDefined]
 ((x)) # E:3: Name "x" is not defined