introduced start and end method call phase explicitly

Author: asmirnov69

pandas_flavor/register.py

@@ -1,19 +1,9 @@
 import ipdb
 from functools import wraps
-import inspect
 from pandas.api.extensions import register_series_accessor, register_dataframe_accessor
 from . import stack_counter
 
-handle_pandas_method_call = None
-
-class LambdaCall:
-    def __init__(self, func):
-        self.func = func
-        self.ret = None
-
-    def __call__(self, *args, **kwargs):
-        self.ret = self.func(*args, **kwargs)
-        return self.ret
+start_method_call = None
 
 def register_dataframe_method(method):
     """Register a function as a method attached to the Pandas DataFrame.
@@ -35,25 +25,21 @@ def __init__(self, pandas_obj):
                 self._obj = pandas_obj
 
             @wraps(method)
-            def __call__(self, *args, **orig_kwargs):
+            def __call__(self, *args, **kwargs):
                 with stack_counter.global_scf.get_sc() as sc:
                     #ipdb.set_trace()
-                    kwargs = {}
-                    for k, arg in orig_kwargs.items():
-                        if method.__name__ == 'assign' and inspect.isfunction(arg):
-                            print('got got that')
-                            #ipdb.set_trace()
-                            kwargs[k] = LambdaCall(arg)
-                        else:
-                            kwargs[k] = arg
-
+                            
+                    method_call_obj = None
+                    if sc.scf.level == 1:
+                        global start_method_call
+                        if start_method_call:
+                            method_call_obj = start_method_call(self._obj, method.__name__, args, kwargs)
+                    
                     ret = method(self._obj, *args, **kwargs)
 
-                    if sc.scf.level == 1:
-                        #print("pf dataframe __call__")
-                        global handle_pandas_method_call
-                        if handle_pandas_method_call:
-                            handle_pandas_method_call(self._obj, method.__name__, args, kwargs, ret)
+                    if method_call_obj:
+                        method_call_obj.handle_end_method_call(ret)
+                    
                     return ret
 
         register_dataframe_accessor(method.__name__)(AccessorMethod)
@@ -76,13 +62,17 @@ def __init__(self, pandas_obj):
             @wraps(method)
             def __call__(self, *args, **kwargs):
                 with stack_counter.global_scf.get_sc() as sc:
-                    ret = method(self._obj, *args, **kwargs)
+                    method_call_obj = None
                     if sc.scf.level <= 2:
-                        #ipdb.set_trace()
-                        print("pf series __call__")
-                        global handle_pandas_method_call
-                        if handle_pandas_method_call:
-                            handle_pandas_method_call(self._obj, method.__name__, args, kwargs, ret)
+                        global start_method_call
+                        if start_method_call:
+                            method_call_obj = start_method_call(self._obj, method.__name__, args, kwargs)
+                            
+                    ret = method(self._obj, *args, **kwargs)
+
+                    if method_call_obj:
+                        method_call_obj.handle_end_method_call(ret)
+                        
                     return ret
 
         register_series_accessor(method.__name__)(AccessorMethod)