Merge pull request #676 from DanielYang59/deprecation-msg

Extend `dev.deprecated` to decorate classes and improve message

monty/dev.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import functools
+import inspect
 import logging
 import os
 import subprocess
@@ -64,9 +65,9 @@ def _is_in_owner_repo() -> bool:
                 owner_repo = (
                     result.stdout.decode("utf-8")
                     .strip()
-                    .lstrip("https://github.com/")  # https clone
-                    .lstrip("git@github.com:")  # ssh clone
-                    .rstrip(".git")  # ssh clone
+                    .lstrip("https://github.com/")  # HTTPS clone
+                    .lstrip("git@github.com:")  # SSH clone
+                    .rstrip(".git")  # SSH clone
                 )
 
                 return owner_repo == os.getenv("GITHUB_REPOSITORY")
@@ -103,27 +104,53 @@ def craft_message(
                 r = replacement.__func__
             else:
                 r = replacement
-            msg += f"; use {r.__name__} in {r.__module__} instead."
+
+            if deadline is None:
+                msg += "; use "  # for better formatting
+            else:
+                msg += "Use "
+            msg += f"{r.__name__} in {r.__module__} instead."
 
         if message:
             msg += "\n" + message
         return msg
 
-    def deprecated_decorator(old: Callable) -> Callable:
+    def deprecated_function_decorator(old: Callable) -> Callable:
         def wrapped(*args, **kwargs):
             msg = craft_message(old, replacement, message, _deadline)
             warnings.warn(msg, category=category, stacklevel=2)
             return old(*args, **kwargs)
 
         return wrapped
 
+    def deprecated_class_decorator(cls: Type) -> Type:
+        original_init = cls.__init__
+
+        def new_init(self, *args, **kwargs):
+            msg = craft_message(cls, replacement, message, _deadline)
+            warnings.warn(msg, category=category, stacklevel=2)
+            original_init(self, *args, **kwargs)
+
+        cls.__init__ = new_init
+        return cls
+
     # Convert deadline to datetime type
     _deadline = datetime(*deadline) if deadline is not None else None
 
     # Raise CI warning after removal deadline
     raise_deadline_warning()
 
-    return deprecated_decorator
+    def decorator(target: Callable) -> Callable:
+        if inspect.isfunction(target):
+            return deprecated_function_decorator(target)
+        elif inspect.isclass(target):
+            return deprecated_class_decorator(target)
+        else:
+            raise TypeError(
+                "The @deprecated decorator can only be applied to classes or functions"
+            )
+
+    return decorator
 
 
 class requires:

tests/test_dev.py

@@ -92,6 +92,23 @@ def classmethod_b(cls):
         with pytest.warns(DeprecationWarning):
             assert TestClass_deprecationwarning().classmethod_b() == "b"
 
+    def test_deprecated_class(self):
+        class TestClassNew:
+            """A dummy class for tests."""
+
+            def method_a(self):
+                pass
+
+        @deprecated(replacement=TestClassNew)
+        class TestClassOld:
+            """A dummy class for tests."""
+
+            def method_b(self):
+                pass
+
+        with pytest.warns(FutureWarning, match="TestClassOld is deprecated"):
+            TestClassOld()
+
     def test_deprecated_deadline(self, monkeypatch):
         with pytest.raises(DeprecationWarning):
             with patch("subprocess.run") as mock_run: