From e9a0ef8310c0b7098a9ff4990ad3d55d4c07159b Mon Sep 17 00:00:00 2001
From: Nick Mills-Barrett <nick@fizzadar.com>
Date: Sat, 9 Mar 2024 11:27:40 +0000
Subject: [PATCH] Add deprecation flags/warning for operations

---
 pyinfra/api/operation.py | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/pyinfra/api/operation.py b/pyinfra/api/operation.py
index 33b3180e3..7e40050fa 100644
--- a/pyinfra/api/operation.py
+++ b/pyinfra/api/operation.py
@@ -164,6 +164,8 @@ def add_op(state: State, op_func, *args, **kwargs):
 def operation(
     is_idempotent: bool = True,
     idempotent_notice: Optional[str] = None,
+    is_deprecated: bool = False,
+    deprecated_for: Optional[str] = None,
     _set_in_op: bool = True,
 ) -> Callable[[Callable[P, Generator]], PyinfraOperation[P]]:
     """
@@ -175,6 +177,8 @@ def operation(
     def decorator(f: Callable[P, Generator]) -> PyinfraOperation[P]:
         f.is_idempotent = is_idempotent  # type: ignore[attr-defined]
         f.idempotent_notice = idempotent_notice  # type: ignore[attr-defined]
+        f.is_deprecated = is_deprecated  # type: ignore[attr-defined]
+        f.deprecated_for = deprecated_for  # type: ignore[attr-defined]
         return _wrap_operation(f, _set_in_op=_set_in_op)
 
     return decorator
@@ -192,6 +196,15 @@ def decorated_func(*args: P.args, **kwargs: P.kwargs) -> OperationMeta:
                 + "function to call the underlying operation."
             )
 
+        if func.is_deprecated:  # type: ignore[attr-defined]
+            if func.deprecated_for:  # type: ignore[attr-defined]
+                logger.warning(
+                    f"The {get_operation_name_from_func(func)} operation is "
+                    + f"deprecated, please use: {func.deprecated_for}",  # type: ignore[attr-defined]
+                )
+            else:
+                logger.warning(f"The {get_operation_name_from_func(func)} operation is deprecated")
+
         # Configure operation
         #
         # Get the meta kwargs (globals that apply to all hosts)
@@ -279,6 +292,15 @@ def command_generator() -> Iterator[PyinfraCommand]:
     return cast(PyinfraOperation[P], decorated_func)
 
 
+def get_operation_name_from_func(func):
+    if func.__module__:
+        module_bits = func.__module__.split(".")
+        module_name = module_bits[-1]
+        return "{0}.{1}".format(module_name, func.__name__)
+    else:
+        return func.__name__
+
+
 def generate_operation_name(func, host, kwargs, global_arguments):
     # Generate an operation name if needed (Module/Operation format)
     name = global_arguments.get("name")
@@ -287,14 +309,7 @@ def generate_operation_name(func, host, kwargs, global_arguments):
         names = {name}
     else:
         add_args = True
-
-        if func.__module__:
-            module_bits = func.__module__.split(".")
-            module_name = module_bits[-1]
-            name = "{0}/{1}".format(module_name.title(), func.__name__.title())
-        else:
-            name = func.__name__
-
+        name = get_operation_name_from_func(func)
         names = {name}
 
     if host.current_deploy_name: