Add more tests

mschoettle · mschoettle · commit 37ee83368c60 · 2024-05-07T15:00:35.000-04:00
diff --git a/mypy/fastparse.py b/mypy/fastparse.py
@@ -986,6 +986,7 @@ def do_func_def(
         # 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
+        # TODO: to go to the end of the function name: n.col_offset + 4 + len(n.name)
         def_end_column: int | None = n.col_offset
 
         returns = n.returns
diff --git a/mypy/messages.py b/mypy/messages.py
@@ -265,7 +265,9 @@ def span_from_context(ctx: Context) -> Iterable[int]:
         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 else end_column
+            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,
diff --git a/test-data/unit/check-columns.test b/test-data/unit/check-columns.test
@@ -178,22 +178,46 @@ 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(x: int, foo: Optional[str]):
+    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 g(x: int, foo: Optional[str] = None):
+    def f_args(x, *args):
         pass
 
-    def h(x, foo = None):
+    def f_kwargs(x, *args, **kwargs):
         pass
+[builtins fixtures/tuple.pyi]
 [out]
-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
+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 testColumnNameIsNotDefined]
 ((x)) # E:3: Name "x" is not defined
