single callback to create method call object which is also with context, getting simpler code

Author: asmirnov69

pandas_flavor/register.py

@@ -3,8 +3,6 @@
 import inspect
 from contextlib import nullcontext
 
-cb_notify_dataframe_method_call = None
-cb_notify_series_method_call = None
 cb_create_call_stack_context_manager = None
 
 def register_dataframe_method(method):
@@ -31,18 +29,15 @@ def __init__(self, pandas_obj):
             @wraps(method)
             def __call__(self, *args, **kwargs):
                 global cb_create_call_stack_context_manager
-                with cb_create_call_stack_context_manager(method.__name__) if cb_create_call_stack_context_manager else nullcontext():
-                    method_call_obj = None
-                    global cb_notify_dataframe_method_call
-                    if cb_notify_dataframe_method_call:
-                        method_call_obj = cb_notify_dataframe_method_call(self._obj, method.__name__, method_signature, args, kwargs)
-                        if method_call_obj:
-                            new_args, new_kwargs = method_call_obj.handle_start_method_call()
-                            args = new_args[1:]; kwargs = new_kwargs
+                method_call_obj = cb_create_call_stack_context_manager(method.__name__) if cb_create_call_stack_context_manager else nullcontext()
+                with method_call_obj:
+                    if not isinstance(method_call_obj, nullcontext):
+                        new_args, new_kwargs = method_call_obj.handle_start_method_call(self._obj, method.__name__, method_signature, args, kwargs)
+                        args = new_args[1:]; kwargs = new_kwargs
 
                     ret = method(self._obj, *args, **kwargs)
 
-                    if method_call_obj:
+                    if not isinstance(method_call_obj, nullcontext):
                         method_call_obj.handle_end_method_call(ret)
 
                     return ret
@@ -69,18 +64,15 @@ def __init__(self, pandas_obj):
             @wraps(method)
             def __call__(self, *args, **kwargs):
                 global cb_create_call_stack_context_manager
-                with cb_create_call_stack_context_manager(method.__name__) if cb_create_call_stack_context_manager else nullcontext():
-                    method_call_obj = None
-                    global cb_notify_series_method_call
-                    if cb_notify_series_method_call:
-                        method_call_obj = cb_notify_series_method_call(self._obj, method.__name__, method_signature, args, kwargs)
-                        if method_call_obj:
-                            new_args, new_kwargs = method_call_obj.handle_start_method_call()
-                            args = new_args[1:]; kwargs = new_kwargs
+                method_call_obj = cb_create_call_stack_context_manager(method.__name__) if cb_create_call_stack_context_manager else nullcontext()
+                with method_call_obj:
+                    if not isinstance(method_call_obj, nullcontext):
+                        new_args, new_kwargs = method_call_obj.handle_start_method_call(self._obj, method.__name__, method_signature, args, kwargs)
+                        args = new_args[1:]; kwargs = new_kwargs
 
                     ret = method(self._obj, *args, **kwargs)
 
-                    if method_call_obj:
+                    if not isinstance(method_call_obj, nullcontext):
                         method_call_obj.handle_end_method_call(ret)
 
                     return ret