REF: pass align_keys to apply (pandas-dev#32846)

Author: jbrockmendel

pandas/core/internals/managers.py

@@ -27,7 +27,7 @@
 )
 from pandas.core.dtypes.concat import concat_compat
 from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries
+from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 import pandas.core.algorithms as algos
@@ -375,7 +375,7 @@ def reduce(self, func, *args, **kwargs):
 
         return res
 
-    def apply(self: T, f, filter=None, **kwargs) -> T:
+    def apply(self: T, f, filter=None, align_keys=None, **kwargs) -> T:
         """
         Iterate over the blocks, collect and create a new BlockManager.
 
@@ -390,6 +390,7 @@ def apply(self: T, f, filter=None, **kwargs) -> T:
         -------
         BlockManager
         """
+        align_keys = align_keys or []
         result_blocks = []
         # fillna: Series/DataFrame is responsible for making sure value is aligned
 
@@ -404,28 +405,14 @@ def apply(self: T, f, filter=None, **kwargs) -> T:
 
         self._consolidate_inplace()
 
+        align_copy = False
         if f == "where":
             align_copy = True
-            if kwargs.get("align", True):
-                align_keys = ["other", "cond"]
-            else:
-                align_keys = ["cond"]
-        elif f == "putmask":
-            align_copy = False
-            if kwargs.get("align", True):
-                align_keys = ["new", "mask"]
-            else:
-                align_keys = ["mask"]
-        else:
-            align_keys = []
 
-        # TODO(EA): may interfere with ExtensionBlock.setitem for blocks
-        # with a .values attribute.
         aligned_args = {
             k: kwargs[k]
             for k in align_keys
-            if not isinstance(kwargs[k], ABCExtensionArray)
-            and hasattr(kwargs[k], "values")
+            if isinstance(kwargs[k], (ABCSeries, ABCDataFrame))
         }
 
         for b in self.blocks:
@@ -561,13 +548,24 @@ def isna(self, func) -> "BlockManager":
         return self.apply("apply", func=func)
 
     def where(self, **kwargs) -> "BlockManager":
-        return self.apply("where", **kwargs)
+        if kwargs.pop("align", True):
+            align_keys = ["other", "cond"]
+        else:
+            align_keys = ["cond"]
+
+        return self.apply("where", align_keys=align_keys, **kwargs)
 
     def setitem(self, indexer, value) -> "BlockManager":
         return self.apply("setitem", indexer=indexer, value=value)
 
     def putmask(self, **kwargs):
-        return self.apply("putmask", **kwargs)
+
+        if kwargs.pop("align", True):
+            align_keys = ["new", "mask"]
+        else:
+            align_keys = ["mask"]
+
+        return self.apply("putmask", align_keys=align_keys, **kwargs)
 
     def diff(self, n: int, axis: int) -> "BlockManager":
         return self.apply("diff", n=n, axis=axis)