REF/ENH: Refactor NDFrame finalization

pandas/core/_meta.py

@@ -0,0 +1,146 @@
+"""
+Metadata propagation through pandas operations.
+
+This module contains the infrastructure for propagating ``NDFrame._metadata``
+through operations. We perform an operation (say :meth:`pandas.Series.copy`) that
+returns an ``NDFrame`` and would like to propagate the metadata (say ``Series.name``)
+from ``self`` to the new ``NDFrame``.
+
+.. note::
+
+   Currently, pandas doesn't provide a clean, documented API on
+
+   * which methods call finalize
+   * the types passed to finalize for each method
+
+   This is a known limitation we would like to address in the future.
+"""
+from collections import defaultdict
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, Union
+
+from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
+
+if TYPE_CHECKING:
+    from pandas.core.generic import NDFrame
+
+dispatch = defaultdict(dict)
+dispatch_method_type = Union[Callable[..., "NDFrame"], str]
+
+
+def key_of(method):
+    if isinstance(method, str):
+        # TODO: figure out if this is OK. May be necessary when we have
+        #   things like pd.merge and DataFrame.merge that hit the same finalize.
+        return method
+    elif method:
+        return method.__module__, method.__name__
+
+
+class PandasMetadata:
+    """
+    Dispatch metadata finalization for pandas metadata.
+
+    Users should instantiate a single `PandasMetadata` instance
+    for their piece of metadata and register finalizers for various
+    pandas methods using :meth:`PandsaMetadata.register`.
+
+    Parameters
+    ----------
+    name : str
+        The name of the attribute being finalized.
+
+    Examples
+    --------
+    >>> maxmeta = PandasMetadata("attr")
+
+    Register a finalizer for a given pandas method:
+
+    >>> @maxmeta.register(pd.concat)
+    ... def _(new, concatenator):
+    ...     new.attr = max(x.attr_meta for x in concatenator.objs)
+
+    >>> pd.DataFrame._metadata = ['attr']
+    >>> x = pd.DataFrame({"x"}); x.attr = 1
+    >>> y = pd.DataFrame({"y"}); y.attr = 2
+    >>> pd.concat([x, y]).attr
+    2
+    """
+
+    def __init__(self, name: str):
+        self.name = name
+
+    def register(self, pandas_method: dispatch_method_type):
+        """
+        A decorator to register a finalizer for a specific pandas method.
+
+        Parameters
+        ----------
+        pandas_method : callable or str
+            A pandas method, like :meth:`pandas.concat`, that this finalizer
+            should be used for. The function being decorated will be called
+            with the relevant arguments (typically the output and the source NDFrame).
+            When `NDFrame.__finalize__` is called as a result of `pandas_method`,
+            the registered finalizer will be called.
+        """
+
+        def decorate(func):
+            # TODO: warn of collisions?
+            dispatch[key_of(pandas_method)][self.name] = func
+
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorate
+
+
+def default_finalizer(new: "NDFrame", other: Any, *, name: str):
+    """
+    The default finalizer when this method, attribute hasn't been overridden.
+
+    This copies the ``_metadata`` attribute from ``other`` to ``self``, modifying
+    ``self`` inplace.
+
+    Parameters
+    ----------
+    new : NDFrame
+        The newly created NDFrame being finalized.
+    other : NDFrame
+        The source NDFrame attributes will be extracted from.
+    """
+    object.__setattr__(new, name, getattr(other, name, None))
+
+
+# ----------------------------------------------------------------------------
+# Pandas Internals.
+
+
+def ndframe_finalize(new: "NDFrame", other: Any, method: dispatch_method_type):
+    """
+    Finalize a new NDFrame.
+
+    The finalizer is looked up from finalizers registered with PandasMetadata.
+    `new` is modified inplace, and nothing is returned.
+
+    Parameters
+    ----------
+    new : NDFrame
+    other : NDFrame
+        Or a list of them? TBD
+    method : callable or str
+    """
+    # To avoid one isinstance per _metadata name, we check up front.
+    # Most of the time `other` is an ndframe, but in some cases (e.g. concat)
+    # it's `_Concatenator` object
+    other_is_ndframe = isinstance(other, (ABCSeries, ABCDataFrame))
+
+    for name in new._metadata:
+        finalizer = dispatch.get(key_of(method), {}).get(name)
+
+        if finalizer:
+            finalizer(new, other)
+        elif other_is_ndframe:
+            default_finalizer(new, other, name=name)

pandas/core/generic.py

@@ -62,6 +62,7 @@
 import pandas as pd
 from pandas._typing import Dtype, FilePathOrBuffer
 from pandas.core import missing, nanops
+from pandas.core._meta import ndframe_finalize
 import pandas.core.algorithms as algos
 from pandas.core.base import PandasObject, SelectionMixin
 import pandas.core.common as com
@@ -5175,9 +5176,7 @@ def __finalize__(self, other, method=None, **kwargs):
             types of propagation actions based on this
 
         """
-        if isinstance(other, NDFrame):
-            for name in self._metadata:
-                object.__setattr__(self, name, getattr(other, name, None))
+        ndframe_finalize(self, other, method)
         return self
 
     def __getattr__(self, name):
@@ -6016,7 +6015,7 @@ def copy(self, deep=True):
         dtype: object
         """
         data = self._data.copy(deep=deep)
-        return self._constructor(data).__finalize__(self)
+        return self._constructor(data).__finalize__(self, NDFrame.copy)
 
     def __copy__(self, deep=True):
         return self.copy(deep=deep)

pandas/tests/generic/test_metadata.py

@@ -0,0 +1,49 @@
+import pytest
+
+import pandas as pd
+from pandas.core._meta import PandasMetadata
+
+mymeta = PandasMetadata("attr")
+
+
+@mymeta.register(pd.core.generic.NDFrame.copy)
+def _(new, other):
+    new.attr = other.attr + 1
+
+
+@mymeta.register("concat")
+def _(new, other):
+    assert isinstance(other, pd.core.reshape.concat._Concatenator)
+    new.attr = sum(x.attr for x in other.objs)
+
+
+@pytest.fixture
+def custom_meta(monkeypatch):
+    original_metadata = []
+
+    for cls in [pd.Series, pd.DataFrame]:
+        original_metadata.append(cls._metadata)
+        custom_metadata = cls._metadata.copy()
+        custom_metadata.append("attr")
+
+        monkeypatch.setattr(cls, "_metadata", custom_metadata)
+
+
+def test_custom_finalizer(custom_meta):
+
+    df = pd.DataFrame({"A": [1, 2]})
+    df.attr = 0
+
+    result = df.copy()
+    assert result.attr == 1
+
+
+def test_concat(custom_meta):
+    df1 = pd.DataFrame({"A": [1, 2]})
+    df1.attr = 2
+
+    df2 = pd.DataFrame({"A": [1, 2]})
+    df2.attr = 3
+
+    result = pd.concat([df1, df2])
+    assert result.attr == 5