Compute end_line and end_column for FuncDef so they don't include the body

shaperilio · shaperilio · commit 10c391bdbb39 · 2024-01-06T12:44:00.000-08:00
diff --git a/mypy/fastparse.py b/mypy/fastparse.py
@@ -985,9 +985,45 @@ 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)
+        # We don't want the end positions for functions to include the entire
+        # body, which is what `ast` gives us; we want to highlight only the
+        # function signature as the error / note.
+
+        # First, check for return typehint
+        ret = n.returns
+        if ret is not None:
+            # There is a return typehint; highlight to the end of it.
+            end_line = ret.end_lineno
+            end_column = ret.end_col_offset
+        elif len(n.body) > 0:
+            # There's no return type, but there is a body (which includes
+            # docstrings).
+            def_line = n.lineno
+            body_line = n.body[0].lineno
+            if def_line == body_line:
+                # Single line definition, e.g. `def foo(): pass`
+                # NOTE: this will not highlight the colon in the nonstandard
+                # case of `def foo():pass`
+                end_line = def_line
+                end_column = n.body[0].col_offset - 1
+            else:
+                # "Proper" body starting on a different line after `:`
+                # We highlight up to the line in which the body starts
+                # but at column 0, effectively ending at the end of the
+                # previous line.
+                # NOTE: this causes funny highlighting in the non-standard
+                # case below:
+                #
+                # def foo(x: int
+                #          ): pass
+                #
+                # The error will show from `d` of `def to `t` of `int`.
+                end_line = n.body[0].lineno
+                end_column = 0
+        else:
+            # Fall back to whole function.
+            end_line = getattr(n, "end_lineno", None)
+            end_column = getattr(n, "end_col_offset", None)
 
         self.class_and_function_stack.pop()
         self.class_and_function_stack.append("F")
