Add wrapper for deprecating functions.

PiperOrigin-RevId: 616781556

chex/__init__.py

@@ -98,6 +98,7 @@
 from chex._src.variants import params_product
 from chex._src.variants import TestCase
 from chex._src.variants import variants
+from chex._src.warnings import warn_deprecated_function
 from chex._src.warnings import warn_keyword_args_only_in_future
 from chex._src.warnings import warn_only_n_pos_args_in_future
 
@@ -188,6 +189,9 @@
     "Shape",
     "TestCase",
     "variants",
+    "warn_deprecated_function",
+    "warn_keyword_args_only_in_future",
+    "warn_only_n_pos_args_in_future",
     "with_jittable_assertions",
 )
 

chex/_src/warnings.py

@@ -55,3 +55,30 @@ def wrapper(*args, **kwargs):
 warn_keyword_args_only_in_future = functools.partial(
     warn_only_n_pos_args_in_future, n=0
 )
+
+
+def warn_deprecated_function(fun, replacement):
+  """A decorator to mark a function definition as deprecated.
+
+  Example usage:
+  @warn_deprecated_function(fun, replacement='g')
+  def f(a, b):
+    return a + b
+
+  Args:
+    fun: the deprecated function.
+    replacement: the name of the function to be used instead.
+
+  Returns:
+    the wrapped function.
+  """
+
+  @functools.wraps(fun)
+  def new_fun(*args, **kwargs):
+    warnings.warn(
+        f'The function {fun.__name__} is deprecated, '
+        f'please use {replacement} instead.',
+        category=DeprecationWarning,
+        stacklevel=2)
+    return fun(*args, **kwargs)
+  return new_fun

chex/_src/warnings_test.py

@@ -26,6 +26,11 @@ def f(a, b, c):
   return a + b + c
 
 
+@functools.partial(warnings.warn_deprecated_function, replacement='h')
+def g(a, b, c):
+  return a + b + c
+
+
 class WarningsTest(absltest.TestCase):
 
   def test_warn_only_n_pos_args_in_future(self):
@@ -34,6 +39,10 @@ def test_warn_only_n_pos_args_in_future(self):
     with self.assertWarns(Warning):
       f(1, 2, c=3)
 
+  def test_warn_deprecated_function(self):
+    with self.assertWarns(Warning):
+      g(1, 2, 3)
+
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   absltest.main()

docs/api.rst

@@ -154,6 +154,7 @@ Warnings
 
 .. currentmodule:: chex
 
+.. autofunction:: warn_deprecated_function
 .. autofunction:: warn_keyword_args_only_in_future
 .. autofunction:: warn_only_n_pos_args_in_future