initial commit

venv/Lib/site-packages/pandas/tests/apply/conftest.py

@@ -0,0 +1,63 @@
+import numpy as np
+import pytest
+
+from pandas import (
+    DataFrame,
+    Series,
+)
+from pandas.api.executors import BaseExecutionEngine
+
+
+class MockExecutionEngine(BaseExecutionEngine):
+    """
+    Execution Engine to test if the execution engine interface receives and
+    uses all parameters provided by the user.
+
+    Making this engine work as the default Python engine by calling it, no extra
+    functionality is implemented here.
+
+    When testing, this will be called when this engine is provided, and then the
+    same pandas.map and pandas.apply function will be called, but without engine,
+    executing the default behavior from the python engine.
+    """
+
+    def map(data, func, args, kwargs, decorator, skip_na):
+        kwargs_to_pass = kwargs if isinstance(data, DataFrame) else {}
+        return data.map(func, na_action="ignore" if skip_na else None, **kwargs_to_pass)
+
+    def apply(data, func, args, kwargs, decorator, axis):
+        if isinstance(data, Series):
+            return data.apply(func, convert_dtype=True, args=args, by_row=False)
+        elif isinstance(data, DataFrame):
+            return data.apply(
+                func,
+                axis=axis,
+                raw=False,
+                result_type=None,
+                args=args,
+                by_row="compat",
+                **kwargs,
+            )
+        else:
+            assert isinstance(data, np.ndarray)
+
+            def wrap_function(func):
+                # https://github.com/numpy/numpy/issues/8352
+                def wrapper(*args, **kwargs):
+                    result = func(*args, **kwargs)
+                    if isinstance(result, str):
+                        result = np.array(result, dtype=object)
+                    return result
+
+                return wrapper
+
+            return np.apply_along_axis(wrap_function(func), axis, data, *args, **kwargs)
+
+
+class MockEngineDecorator:
+    __pandas_udf__ = MockExecutionEngine
+
+
+@pytest.fixture(params=[None, MockEngineDecorator])
+def engine(request):
+    return request.param