Merge pull request #28 from pyjanitor-devs/pyjviz-callbacks

addition of method_call_ctx_factory

Author: ericmjl

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
   rev: 1.5.0
   hooks:
   - id: interrogate
-    args: [-c, pyproject.toml]
+    args: [-c, pyproject.toml, -vv]
 - repo: https://github.com/terrencepreilly/darglint
   rev: v1.8.1
   hooks:

README.md

@@ -147,11 +147,11 @@ If you have a feature request, please open an issue or submit a PR!
 
 ## TL;DR
 
-Pandas 0.23 introduced a simpler API for [extending Pandas](https://pandas.pydata.org/pandas-docs/stable/development/extending.html#extending-pandas). This API provided two key decorators, `register_dataframe_accessor` and `register_series_accessor`, that enable users to register **accessors** with Pandas DataFrames and Series. 
+Pandas 0.23 introduced a simpler API for [extending Pandas](https://pandas.pydata.org/pandas-docs/stable/development/extending.html#extending-pandas). This API provided two key decorators, `register_dataframe_accessor` and `register_series_accessor`, that enable users to register **accessors** with Pandas DataFrames and Series.
 
-Pandas Flavor originated as a library to backport these decorators to older versions of Pandas (<0.23). While doing the backporting, it became clear that registering **methods** directly to Pandas objects might be a desired feature as well.[*](#footnote) 
+Pandas Flavor originated as a library to backport these decorators to older versions of Pandas (<0.23). While doing the backporting, it became clear that registering **methods** directly to Pandas objects might be a desired feature as well.[*](#footnote)
 
-<a name="footnote">*</a>*It is likely that Pandas deliberately chose not implement to this feature. If everyone starts monkeypatching DataFrames with their custom methods, it could lead to confusion in the Pandas community. The preferred Pandas approach is to namespace your methods by registering an accessor that contains your custom methods.* 
+<a name="footnote">*</a>*It is likely that Pandas deliberately chose not implement to this feature. If everyone starts monkeypatching DataFrames with their custom methods, it could lead to confusion in the Pandas community. The preferred Pandas approach is to namespace your methods by registering an accessor that contains your custom methods.*
 
 **So how does method registration work?**
 

pandas_flavor/__version__.py

@@ -1,2 +1,2 @@
 """Version number."""
-__version__ = "0.3.0"
+__version__ = "0.4.0"

pandas_flavor/register.py

@@ -1,29 +1,160 @@
 from functools import wraps
-from pandas.api.extensions import register_series_accessor, register_dataframe_accessor
+from pandas.api.extensions import (
+    register_series_accessor,
+    register_dataframe_accessor,
+)
+import inspect
+
+method_call_ctx_factory = None
+
+
+def handle_pandas_extension_call(method, method_signature, obj, args, kwargs):
+    """Handle pandas extension call.
+
+    This function is called when the user calls
+    a pandas DataFrame object's method.
+    The pandas extension mechanism passes args and kwargs
+    of the original method call as it is applied to obj.
+
+    Our implementation uses the global variable `method_call_ctx_factory`.
+
+    `method_call_ctx_factory` can be either None or an abstract class.
+
+    When `method_call_ctx_factory` is None,
+    the implementation calls the registered method
+    with unmodified args and kwargs and returns underlying method result.
+
+    When `method_call_ctx_factory` is not None,
+    `method_call_ctx_factory` is expected to refer to
+    the function to create the context object.
+    The context object will be used to process
+    inputs and outputs of `method` calls.
+    It is also possible that
+    the context object method `handle_start_method_call`
+    will modify original args and kwargs before `method` call.
+
+    `method_call_ctx_factory` is a function
+    that should have the following signature:
+
+     `f(method_name: str, args: list, kwargs: dict) -> MethodCallCtx`
+
+
+    MethodCallCtx is an abstract class:
+    class MethodCallCtx(abc.ABC):
+
+        @abstractmethod
+        def __enter__(self) -> None:
+            raise NotImplemented
+
+        @abstractmethod
+        def __exit__(self, exc_type, exc_value, traceback) -> None:
+            raise NotImplemented
+
+        @abstractmethod
+        def handle_start_method_call(self, method_name: str, method_signature: inspect.Signature, method_args: list, method_kwargs: dict) -> tuple(list, dict):
+            raise NotImplemented
+
+        @abstractmethod
+        def handle_end_method_call(self, ret: object) -> None:
+            raise NotImplemented
+
+
+    Args:
+        method (callable): method object as registered by decorator
+            register_dataframe_method (or register_series_method)
+        method_signature: signature of method as returned by inspect.signature
+        obj: Dataframe or Series
+        args: The arguments to pass to the registered method.
+        kwargs: The keyword arguments to pass to the registered method.
+
+    Returns:
+        object`: The result of calling of the method.
+    """  # noqa: E501
+
+    global method_call_ctx_factory
+    with method_call_ctx_factory(
+        method.__name__, args, kwargs
+    ) as method_call_ctx:
+        if method_call_ctx is None:  # nullcontext __enter__ returns None
+            ret = method(obj, *args, **kwargs)
+        else:
+            all_args = tuple([obj] + list(args))
+            (new_args, new_kwargs,) = method_call_ctx.handle_start_method_call(
+                method.__name__, method_signature, all_args, kwargs
+            )
+            args = new_args[1:]
+            kwargs = new_kwargs
+
+            ret = method(obj, *args, **kwargs)
+
+            method_call_ctx.handle_end_method_call(ret)
+
+        return ret
 
 
 def register_dataframe_method(method):
     """Register a function as a method attached to the Pandas DataFrame.
 
-    Example
-    -------
-
-    .. code-block:: python
+    Example:
 
         @register_dataframe_method
         def print_column(df, col):
             '''Print the dataframe column given'''
             print(df[col])
+
+    Args:
+        method (callable): callable to register as a dataframe method.
+
+    Returns:
+        callable: The original method.
     """
 
+    method_signature = inspect.signature(method)
+
     def inner(*args, **kwargs):
+        """Inner function to register the method.
+
+        This function is called when the user
+        decorates a function with register_dataframe_method.
+
+        Args:
+            *args: The arguments to pass to the registered method.
+            **kwargs: The keyword arguments to pass to the registered method.
+
+        Returns:
+            method: The original method.
+        """
+
         class AccessorMethod(object):
+            """DataFrame Accessor method class."""
+
             def __init__(self, pandas_obj):
+                """Initialize the accessor method class.
+
+                Args:
+                    pandas_obj (pandas.DataFrame): The pandas DataFrame object.
+                """
                 self._obj = pandas_obj
 
             @wraps(method)
             def __call__(self, *args, **kwargs):
-                return method(self._obj, *args, **kwargs)
+                """Call the accessor method.
+
+                Args:
+                    *args: The arguments to pass to the registered method.
+                    **kwargs: The keyword arguments to pass
+                        to the registered method.
+
+                Returns:
+                    object: The result of calling of the method.
+                """
+                global method_call_ctx_factory
+                if method_call_ctx_factory is None:
+                    return method(self._obj, *args, **kwargs)
+
+                return handle_pandas_extension_call(
+                    method, method_signature, self._obj, args, kwargs
+                )
 
         register_dataframe_accessor(method.__name__)(AccessorMethod)
 
@@ -33,18 +164,50 @@ def __call__(self, *args, **kwargs):
 
 
 def register_series_method(method):
-    """Register a function as a method attached to the Pandas Series."""
+    """Register a function as a method attached to the Pandas Series.
+
+    Args:
+        method (callable): callable to register as a series method.
+
+    Returns:
+        callable: The original method.
+    """
+
+    method_signature = inspect.signature(method)
 
     def inner(*args, **kwargs):
         class AccessorMethod(object):
+            """Series Accessor method class."""
+
             __doc__ = method.__doc__
 
             def __init__(self, pandas_obj):
+                """Initialize the accessor method class.
+
+                Args:
+                    pandas_obj (pandas.Series): The pandas Series object.
+                """
                 self._obj = pandas_obj
 
             @wraps(method)
             def __call__(self, *args, **kwargs):
-                return method(self._obj, *args, **kwargs)
+                """Call the accessor method.
+
+                Args:
+                    *args: The arguments to pass to the registered method.
+                    **kwargs: The keyword arguments to pass
+                        to the registered method.
+
+                Returns:
+                    object: The result of calling of the method.
+                """
+                global method_call_ctx_factory
+                if method_call_ctx_factory is None:
+                    return method(self._obj, *args, **kwargs)
+
+                return handle_pandas_extension_call(
+                    method, method_signature, self._obj, args, kwargs
+                )
 
         register_series_accessor(method.__name__)(AccessorMethod)
 

pandas_flavor/xarray.py

@@ -1,3 +1,4 @@
+"""XArray support for pandas_flavor."""
 from xarray import register_dataarray_accessor, register_dataset_accessor
 from functools import wraps
 
@@ -6,29 +7,67 @@ def make_accessor_wrapper(method):
     """
     Makes an XArray-compatible accessor to wrap a method to be added to an
     xr.DataArray, xr.Dataset, or both.
-    :param method: A method which takes an XArray object and needed parameters.
-    :return: The result of calling ``method``.
+
+    Args:
+        method: A method which takes an XArray object and needed parameters.
+
+    Returns:
+        The result of calling ``method``.
     """
 
     class XRAccessor:
+        """XArray accessor for a method."""
+
         def __init__(self, xr_obj):
+            """Initialize the accessor.
+
+            Args:
+                xr_obj: The XArray object to which the accessor is attached.
+            """
             self._xr_obj = xr_obj
 
         @wraps(method)
         def __call__(self, *args, **kwargs):
+            """Call the method.
+
+            Args:
+                *args: Positional arguments to pass to the method.
+                **kwargs: Keyword arguments to pass to the method.
+
+            Returns:
+                The result of calling ``method``.
+            """
+
             return method(self._xr_obj, *args, **kwargs)
 
     return XRAccessor
 
 
 def register_xarray_dataarray_method(method: callable):
+    """Register a method on an XArray DataArray object.
+
+    Args:
+        method: A method which takes an XArray object and needed parameters.
+
+    Returns:
+        The method.
+    """
     accessor_wrapper = make_accessor_wrapper(method)
     register_dataarray_accessor(method.__name__)(accessor_wrapper)
 
     return method
 
 
 def register_xarray_dataset_method(method: callable):
+    """Register a method on an XArray Dataset object.
+
+    Args:
+        method: A method which takes an XArray object and needed parameters.
+
+    Returns:
+        The method.
+    """
+
     accessor_wrapper = make_accessor_wrapper(method)
     register_dataset_accessor(method.__name__)(accessor_wrapper)
 

setup.py

@@ -24,12 +24,14 @@
 # The rest you shouldn't have to touch too much :)
 # ------------------------------------------------
 # Except, perhaps the License and Trove Classifiers!
-# If you do change the License, remember to change the Trove Classifier for that!
+# If you do change the License,
+# remember to change the Trove Classifier for that!
 
 here = os.path.abspath(os.path.dirname(__file__))
 
 # Import the README and use it as the long-description.
-# Note: this will only work if 'README.rst' is present in your MANIFEST.in file!
+# Note: this will only work if 'README.rst'
+# is present in your MANIFEST.in file!
 with io.open(os.path.join(here, "README.md"), encoding="utf-8") as f:
     long_description = "\n" + f.read()
 
@@ -46,8 +48,13 @@ class UploadCommand(Command):
     user_options = []
 
     @staticmethod
-    def status(s):
-        """Prints things in bold."""
+    def status(s: str):
+        """Prints things in bold.
+
+        Args:
+            s: The string to print.
+
+        """
         print("\033[1m{0}\033[0m".format(s))
 
     def initialize_options(self):