test: Make skip_requires_pyarrow compatible w/ pytest.param (#3772)

Author: dangotbanned

tests/__init__.py

@@ -5,21 +5,22 @@
 import sys
 from importlib.util import find_spec
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, overload
 
 import pytest
 
 from tests import examples_arguments_syntax, examples_methods_syntax
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection, Iterator, Mapping
+    from collections.abc import Collection, Iterator, Mapping
     from re import Pattern
 
     if sys.version_info >= (3, 11):
         from typing import TypeAlias
     else:
         from typing_extensions import TypeAlias
     from _pytest.mark import ParameterSet
+    from _pytest.mark.structures import Markable
 
     MarksType: TypeAlias = (
         "pytest.MarkDecorator | Collection[pytest.MarkDecorator | pytest.Mark]"
@@ -96,9 +97,21 @@ def windows_has_tzdata() -> bool:
 """
 
 
+@overload
 def skip_requires_pyarrow(
-    fn: Callable[..., Any] | None = None, /, *, requires_tzdata: bool = False
-) -> Callable[..., Any]:
+    fn: None = ..., /, *, requires_tzdata: bool = ...
+) -> pytest.MarkDecorator: ...
+
+
+@overload
+def skip_requires_pyarrow(
+    fn: Markable, /, *, requires_tzdata: bool = ...
+) -> Markable: ...
+
+
+def skip_requires_pyarrow(
+    fn: Markable | None = None, /, *, requires_tzdata: bool = False
+) -> pytest.MarkDecorator | Markable:
     """
     ``pytest.mark.skipif`` decorator.
 
@@ -109,7 +122,7 @@ def skip_requires_pyarrow(
     https://github.com/vega/altair/issues/3050
 
     .. _pyarrow:
-    https://pypi.org/project/pyarrow/
+        https://pypi.org/project/pyarrow/
     """
     composed = pytest.mark.skipif(
         find_spec("pyarrow") is None, reason="`pyarrow` not installed."
@@ -120,13 +133,7 @@ def skip_requires_pyarrow(
             reason="Timezone database is not installed on Windows",
         )(composed)
 
-    def wrap(test_fn: Callable[..., Any], /) -> Callable[..., Any]:
-        return composed(test_fn)
-
-    if fn is None:
-        return wrap
-    else:
-        return wrap(fn)
+    return composed if fn is None else composed(fn)
 
 
 def id_func_str_only(val) -> str: