d33bs
diff --git a/‎src/google/adk/utils/feature_decorator.py
Lines changed: 93 additions & 62 deletions b/‎src/google/adk/utils/feature_decorator.py
Lines changed: 93 additions & 62 deletions
diff --git a/‎tests/unittests/utils/test_feature_decorator.py
Lines changed: 91 additions & 0 deletions b/‎tests/unittests/utils/test_feature_decorator.py
Lines changed: 91 additions & 0 deletions
@@ -34,70 +34,90 @@ 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}"
-
       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"
+) -> 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)
+
+    # 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"
         )
 
-    return decorator
+        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"
+        )
 
-  return decorator_factory
+        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(
@@ -137,8 +157,19 @@ def my_wip_function():
 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
 ```
 """
@@ -30,6 +30,31 @@ def run(self):
     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<
3D7A
/span> "This feature is experimental and may change or be removed" in str(
+        w[0].message
+    )