diff --git a/src/google/adk/utils/feature_decorator.py b/src/google/adk/utils/feature_decorator.py index e3be0a65..d597063a 100644 --- a/src/google/adk/utils/feature_decorator.py +++ b/src/google/adk/utils/feature_decorator.py @@ -34,72 +34,92 @@ def _make_feature_decorator( default_message: str, block_usage: bool = False, bypass_env_var: Optional[str] = None, -) -> Callable[[str], Callable[[T], T]]: - def decorator_factory(message: str = default_message) -> Callable[[T], T]: - def decorator(obj: T) -> T: - obj_name = getattr(obj, "__name__", type(obj).__name__) - msg = f"[{label.upper()}] {obj_name}: {message}" +) -> Callable: + def decorator_factory(message_or_obj=None): + # Case 1: Used as @decorator without parentheses + # message_or_obj is the decorated class/function + if message_or_obj is not None and ( + isinstance(message_or_obj, type) or callable(message_or_obj) + ): + return _create_decorator( + default_message, label, block_usage, bypass_env_var + )(message_or_obj) - if isinstance(obj, type): # decorating a class - orig_init = obj.__init__ - - @functools.wraps(orig_init) - def new_init(self, *args, **kwargs): - # Load .env file if dotenv is available - load_dotenv() - - # Check if usage should be bypassed via environment variable at call time - should_bypass = ( - bypass_env_var is not None - and os.environ.get(bypass_env_var, "").lower() == "true" - ) - - if should_bypass: - # Bypass completely - no warning, no error - pass - elif block_usage: - raise RuntimeError(msg) - else: - warnings.warn(msg, category=UserWarning, stacklevel=2) - return orig_init(self, *args, **kwargs) - - obj.__init__ = new_init # type: ignore[attr-defined] - return cast(T, obj) - - elif callable(obj): # decorating a function or method - - @functools.wraps(obj) - def wrapper(*args, **kwargs): - # Load .env file if dotenv is available - load_dotenv() - - # Check if usage should be bypassed via environment variable at call time - should_bypass = ( - bypass_env_var is not None - and os.environ.get(bypass_env_var, "").lower() == "true" - ) - - if should_bypass: - # Bypass completely - no warning, no error - pass - elif block_usage: - raise RuntimeError(msg) - else: - warnings.warn(msg, category=UserWarning, stacklevel=2) - return obj(*args, **kwargs) - - return cast(T, wrapper) - - else: - raise TypeError( - f"@{label} can only be applied to classes or callable objects" - ) - - return decorator + # Case 2: Used as @decorator() with or without message + # message_or_obj is either None or a string message + message = ( + message_or_obj if isinstance(message_or_obj, str) else default_message + ) + return _create_decorator(message, label, block_usage, bypass_env_var) return decorator_factory +def _create_decorator( + message: str, label: str, block_usage: bool, bypass_env_var: Optional[str] +) -> Callable[[T], T]: + def decorator(obj: T) -> T: + obj_name = getattr(obj, "__name__", type(obj).__name__) + msg = f"[{label.upper()}] {obj_name}: {message}" + + if isinstance(obj, type): # decorating a class + orig_init = obj.__init__ + + @functools.wraps(orig_init) + def new_init(self, *args, **kwargs): + # Load .env file if dotenv is available + load_dotenv() + + # Check if usage should be bypassed via environment variable at call time + should_bypass = ( + bypass_env_var is not None + and os.environ.get(bypass_env_var, "").lower() == "true" + ) + + if should_bypass: + # Bypass completely - no warning, no error + pass + elif block_usage: + raise RuntimeError(msg) + else: + warnings.warn(msg, category=UserWarning, stacklevel=2) + return orig_init(self, *args, **kwargs) + + obj.__init__ = new_init # type: ignore[attr-defined] + return cast(T, obj) + + elif callable(obj): # decorating a function or method + + @functools.wraps(obj) + def wrapper(*args, **kwargs): + # Load .env file if dotenv is available + load_dotenv() + + # Check if usage should be bypassed via environment variable at call time + should_bypass = ( + bypass_env_var is not None + and os.environ.get(bypass_env_var, "").lower() == "true" + ) + + if should_bypass: + # Bypass completely - no warning, no error + pass + elif block_usage: + raise RuntimeError(msg) + else: + warnings.warn(msg, category=UserWarning, stacklevel=2) + return obj(*args, **kwargs) + + return cast(T, wrapper) + + else: + raise TypeError( + f"@{label} can only be applied to classes or callable objects" + ) + + return decorator + + working_in_progress = _make_feature_decorator( label="WIP", default_message=( @@ -137,8 +157,19 @@ experimental = _make_feature_decorator( Sample usage: ``` -@experimental("This API may have breaking change in the future.") +# Use with default message +@experimental class ExperimentalClass: pass + +# Use with custom message +@experimental("This API may have breaking change in the future.") +class CustomExperimentalClass: + pass + +# Use with empty parentheses (same as default message) +@experimental() +def experimental_function(): + pass ``` """ diff --git a/tests/unittests/utils/test_feature_decorator.py b/tests/unittests/utils/test_feature_decorator.py index eb700ea6..e2f16446 100644 --- a/tests/unittests/utils/test_feature_decorator.py +++ b/tests/unittests/utils/test_feature_decorator.py @@ -30,6 +30,31 @@ class ExperimentalClass: return "running experimental" +# Test classes/functions for new usage patterns +@experimental +class ExperimentalClassNoParens: + + def run(self): + return "running experimental without parens" + + +@experimental() +class ExperimentalClassEmptyParens: + + def run(self): + return "running experimental with empty parens" + + +@experimental +def experimental_fn_no_parens(): + return "executing without parens" + + +@experimental() +def experimental_fn_empty_parens(): + return "executing with empty parens" + + def test_working_in_progress_class_raises_error(): """Test that WIP class raises RuntimeError by default.""" # Ensure environment variable is not set @@ -208,3 +233,69 @@ def test_experimental_class_warns(): assert issubclass(w[0].category, UserWarning) assert "[EXPERIMENTAL] ExperimentalClass:" in str(w[0].message) assert "class may change" in str(w[0].message) + + +def test_experimental_class_no_parens_warns(): + """Test that experimental class without parentheses shows default warning.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + exp_class = ExperimentalClassNoParens() + result = exp_class.run() + + assert result == "running experimental without parens" + assert len(w) == 1 + assert issubclass(w[0].category, UserWarning) + assert "[EXPERIMENTAL] ExperimentalClassNoParens:" in str(w[0].message) + assert "This feature is experimental and may change or be removed" in str( + w[0].message + ) + + +def test_experimental_class_empty_parens_warns(): + """Test that experimental class with empty parentheses shows default warning.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + exp_class = ExperimentalClassEmptyParens() + result = exp_class.run() + + assert result == "running experimental with empty parens" + assert len(w) == 1 + assert issubclass(w[0].category, UserWarning) + assert "[EXPERIMENTAL] ExperimentalClassEmptyParens:" in str(w[0].message) + assert "This feature is experimental and may change or be removed" in str( + w[0].message + ) + + +def test_experimental_function_no_parens_warns(): + """Test that experimental function without parentheses shows default warning.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + result = experimental_fn_no_parens() + + assert result == "executing without parens" + assert len(w) == 1 + assert issubclass(w[0].category, UserWarning) + assert "[EXPERIMENTAL] experimental_fn_no_parens:" in str(w[0].message) + assert "This feature is experimental and may change or be removed" in str( + w[0].message + ) + + +def test_experimental_function_empty_parens_warns(): + """Test that experimental function with empty parentheses shows default warning.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + result = experimental_fn_empty_parens() + + assert result == "executing with empty parens" + assert len(w) == 1 + assert issubclass(w[0].category, UserWarning) + assert "[EXPERIMENTAL] experimental_fn_empty_parens:" in str(w[0].message) + assert "This feature is experimental and may change or be removed" in str( + w[0].message + )