From a9f9b15624a4683ba869850ade18ece27cf961b4 Mon Sep 17 00:00:00 2001
From: Philipp A <flying-sheep@web.de>
Date: Thu, 7 Dec 2023 16:27:18 +0100
Subject: [PATCH] Allow customizing warning category (#17)

---
 src/legacy_api_wrap/__init__.py | 17 ++++++++++++++---
 tests/test_basic.py             |  9 +++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/legacy_api_wrap/__init__.py b/src/legacy_api_wrap/__init__.py
index 3b75edd..5ac3fda 100644
--- a/src/legacy_api_wrap/__init__.py
+++ b/src/legacy_api_wrap/__init__.py
@@ -29,7 +29,11 @@
 
 # The actual returned Callable of course accepts more positional parameters,
 # but we want the type to lie so end users don’t rely on the deprecated API.
-def legacy_api(*old_positionals: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
+def legacy_api(
+    *old_positionals: str,
+    category: type[Warning] = DeprecationWarning,
+    stacklevel: int = 2,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """Legacy API wrapper.
 
     You want to change the API of a function:
@@ -54,6 +58,13 @@ def legacy_api(*old_positionals: str) -> Callable[[Callable[P, R]], Callable[P,
     ----------
     old_positionals
         The positional parameter names that the old function had after the new function’s ``*``.
+    category
+        The warning class to use for the deprecation.
+        Typically, you want to use ``DeprecationWarning``, ``PendingDeprecationWarning``,
+        ``FutureWarning``, or a custom subclass of those.
+    stacklevel
+        The stacklevel to use for the deprecation warning.
+        By default, the first stack frame is the call site of the wrapped function.
     """
 
     def wrapper(fn: Callable[P, R]) -> Callable[P, R]:
@@ -82,8 +93,8 @@ def fn_compatible(*args_all: P.args, **kw: P.kwargs) -> R:
                 f"The specified parameters {old_positionals[:len(args_rest)]!r} are "
                 "no longer positional. "
                 f"Please specify them like `{old_positionals[0]}={args_rest[0]!r}`",
-                DeprecationWarning,
-                stacklevel=2,
+                category=category,
+                stacklevel=stacklevel,
             )
             kw_new: P.kwargs = {**kw, **dict(zip(old_positionals, args_rest))}
 
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 0d18e40..ffac046 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -41,3 +41,12 @@ def test_too_many_args() -> None:
         match=r"new\(\) takes from 1 to 4 parameters, but 5 were given\.",
     ):
         new(1, 2, 3, 4, 5)  # type: ignore[misc]
+
+
+def test_customize() -> None:
+    @legacy_api("a", category=FutureWarning)
+    def new(*, a: int) -> int:
+        return a
+
+    with pytest.raises(FutureWarning):
+        new(1)  # type: ignore[misc]