From dbde2d46c9c4439cd8155231e2ecaddf36467356 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Wed, 10 Sep 2025 00:16:21 +0000 Subject: [PATCH 1/2] chore: allow method_logger to directly decorate functions with params in parentheses --- bigframes/core/log_adapter.py | 87 ++++++++++++++++------------- tests/unit/core/test_log_adapter.py | 12 ++++ 2 files changed, 60 insertions(+), 39 deletions(-) diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py index 6021c7075a..4b4afa271a 100644 --- a/bigframes/core/log_adapter.py +++ b/bigframes/core/log_adapter.py @@ -149,49 +149,58 @@ def wrap(cls): return wrap(decorated_cls) -def method_logger(method, /, *, custom_base_name: Optional[str] = None): +def method_logger(method = None, /, *, custom_base_name: Optional[str] = None): """Decorator that adds logging functionality to a method.""" - @functools.wraps(method) - def wrapper(*args, **kwargs): - api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__) - if custom_base_name is None: - qualname_parts = getattr(method, "__qualname__", method.__name__).split(".") - class_name = qualname_parts[-2] if len(qualname_parts) > 1 else "" - base_name = ( - class_name if class_name else "_".join(method.__module__.split(".")[1:]) - ) - else: - base_name = custom_base_name - - full_method_name = f"{base_name.lower()}-{api_method_name}" - # Track directly called methods - if len(_call_stack) == 0: - add_api_method(full_method_name) - - _call_stack.append(full_method_name) - - try: - return method(*args, **kwargs) - except (NotImplementedError, TypeError) as e: - # Log method parameters that are implemented in pandas but either missing (TypeError) - # or not fully supported (NotImplementedError) in BigFrames. - # Logging is currently supported only when we can access the bqclient through - # _block.session.bqclient. - if len(_call_stack) == 1: - submit_pandas_labels( - _get_bq_client(*args, **kwargs), - base_name, - api_method_name, - args, - kwargs, - task=PANDAS_PARAM_TRACKING_TASK, + def outer_wrapper(method): + + @functools.wraps(method) + def wrapper(*args, **kwargs): + api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__) + if custom_base_name is None: + qualname_parts = getattr(method, "__qualname__", method.__name__).split(".") + class_name = qualname_parts[-2] if len(qualname_parts) > 1 else "" + base_name = ( + class_name if class_name else "_".join(method.__module__.split(".")[1:]) ) - raise e - finally: - _call_stack.pop() + else: + base_name = custom_base_name - return wrapper + full_method_name = f"{base_name.lower()}-{api_method_name}" + # Track directly called methods + if len(_call_stack) == 0: + add_api_method(full_method_name) + + _call_stack.append(full_method_name) + + try: + return method(*args, **kwargs) + except (NotImplementedError, TypeError) as e: + # Log method parameters that are implemented in pandas but either missing (TypeError) + # or not fully supported (NotImplementedError) in BigFrames. + # Logging is currently supported only when we can access the bqclient through + # _block.session.bqclient. + if len(_call_stack) == 1: + submit_pandas_labels( + _get_bq_client(*args, **kwargs), + base_name, + api_method_name, + args, + kwargs, + task=PANDAS_PARAM_TRACKING_TASK, + ) + raise e + finally: + _call_stack.pop() + + return wrapper + + if method is None: + # Called with parentheses + return outer_wrapper + + # Called without parentheses + return outer_wrapper(method) def property_logger(prop): diff --git a/tests/unit/core/test_log_adapter.py b/tests/unit/core/test_log_adapter.py index eba015dd9d..3f46fccb25 100644 --- a/tests/unit/core/test_log_adapter.py +++ b/tests/unit/core/test_log_adapter.py @@ -101,6 +101,18 @@ def test_method_logging_with_custom_base_name(test_method_w_custom_base): assert "pandas-method1" in api_methods +def test_method_logging_with_custom_base__logger_as_decorator(): + + @log_adapter.method_logger(custom_base_name="pandas") + def my_method(): + pass + + my_method() + + api_methods = log_adapter.get_and_reset_api_methods() + assert "pandas-my_method" in api_methods + + def test_property_logging(test_instance): test_instance.my_field From 342f3c97316b3e18eb4dafd79a34f0059771d4e5 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Wed, 10 Sep 2025 00:19:59 +0000 Subject: [PATCH 2/2] fix format --- bigframes/core/log_adapter.py | 15 +++++++++------ tests/unit/core/test_log_adapter.py | 1 - 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py index 4b4afa271a..3ec1e86dc7 100644 --- a/bigframes/core/log_adapter.py +++ b/bigframes/core/log_adapter.py @@ -149,19 +149,22 @@ def wrap(cls): return wrap(decorated_cls) -def method_logger(method = None, /, *, custom_base_name: Optional[str] = None): +def method_logger(method=None, /, *, custom_base_name: Optional[str] = None): """Decorator that adds logging functionality to a method.""" def outer_wrapper(method): - @functools.wraps(method) def wrapper(*args, **kwargs): api_method_name = getattr(method, LOG_OVERRIDE_NAME, method.__name__) if custom_base_name is None: - qualname_parts = getattr(method, "__qualname__", method.__name__).split(".") + qualname_parts = getattr(method, "__qualname__", method.__name__).split( + "." + ) class_name = qualname_parts[-2] if len(qualname_parts) > 1 else "" base_name = ( - class_name if class_name else "_".join(method.__module__.split(".")[1:]) + class_name + if class_name + else "_".join(method.__module__.split(".")[1:]) ) else: base_name = custom_base_name @@ -194,13 +197,13 @@ def wrapper(*args, **kwargs): _call_stack.pop() return wrapper - + if method is None: # Called with parentheses return outer_wrapper # Called without parentheses - return outer_wrapper(method) + return outer_wrapper(method) def property_logger(prop): diff --git a/tests/unit/core/test_log_adapter.py b/tests/unit/core/test_log_adapter.py index 3f46fccb25..c236bb6886 100644 --- a/tests/unit/core/test_log_adapter.py +++ b/tests/unit/core/test_log_adapter.py @@ -102,7 +102,6 @@ def test_method_logging_with_custom_base_name(test_method_w_custom_base): def test_method_logging_with_custom_base__logger_as_decorator(): - @log_adapter.method_logger(custom_base_name="pandas") def my_method(): pass