creating function for create_test_df_in_defined_mode

Signed-off-by: arunjose696 <arunjose696@gmail.com>

modin/core/storage_formats/pandas/query_compiler_caster.py

@@ -52,7 +52,7 @@ def __init_subclass__(
         **kwargs : Additional keyword arguments
         """
         super().__init_subclass__(**kwargs)
-        apply_argument_cast()(cls)
+        apply_argument_cast(cls)
 
 
 def cast_nested_args_to_current_qc_type(arguments, current_qc):
@@ -100,61 +100,60 @@ def cast_arg_to_current_qc(arg):
     return arguments
 
 
-def apply_argument_cast():
+def apply_argument_cast(obj: Fn) -> Fn:
     """
-    Cast any of args that is a query compiler to the type of left query compiler.
+    Cast all arguments that are query compilers to the current query compiler.
+
+    Parameters
+    ----------
+    obj : function
 
     Returns
     -------
-    func
-        A decorator function.
+    function
+        Returns decorated function which does argument casting.
     """
+    if isinstance(obj, type):
+        all_attrs = dict(inspect.getmembers(obj))
+        all_attrs.pop("__abstractmethods__")
+
+        # This is required because inspect converts class methods to member functions
+        current_class_attrs = vars(obj)
+        for key in current_class_attrs:
+            all_attrs[key] = current_class_attrs[key]
+
+        for attr_name, attr_value in all_attrs.items():
+            if isinstance(
+                attr_value, (FunctionType, MethodType, classmethod, staticmethod)
+            ):
+                wrapped = apply_argument_cast(attr_value)
+                setattr(obj, attr_name, wrapped)
+        return obj  # type: ignore [return-value]
+    elif isinstance(obj, classmethod):
+        return classmethod(apply_argument_cast(obj.__func__))  # type: ignore [return-value, arg-type]
+    elif isinstance(obj, staticmethod):
+        return staticmethod(apply_argument_cast(obj.__func__))
+
+    @functools.wraps(obj)
+    def cast_args(*args: Tuple, **kwargs: Dict) -> Any:
+        """
+        Add casting for query compiler arguments.
+
+        Parameters
+        ----------
+        *args : tuple
+            The function arguments.
+        **kwargs : dict
+            The function keyword arguments.
+
+        Returns
+        -------
+        Any
+        """
+        current_qc = args[0]
+        if isinstance(current_qc, BaseQueryCompiler):
+            kwargs = cast_nested_args_to_current_qc_type(kwargs, current_qc)
+            args = cast_nested_args_to_current_qc_type(args, current_qc)
+        return obj(*args, **kwargs)
 
-    def decorator(obj: Fn) -> Fn:
-        """Cast all arguments that are query compilers to the current query compiler."""
-        if isinstance(obj, type):
-            all_attrs = dict(inspect.getmembers(obj))
-            all_attrs.pop("__abstractmethods__")
-
-            # This is required because inspect converts class methods to member functions
-            current_class_attrs = vars(obj)
-            for key in current_class_attrs:
-                all_attrs[key] = current_class_attrs[key]
-
-            for attr_name, attr_value in all_attrs.items():
-                if isinstance(
-                    attr_value, (FunctionType, MethodType, classmethod, staticmethod)
-                ):
-                    wrapped = apply_argument_cast()(attr_value)
-                    setattr(obj, attr_name, wrapped)
-            return obj  # type: ignore [return-value]
-        elif isinstance(obj, classmethod):
-            return classmethod(decorator(obj.__func__))  # type: ignore [return-value, arg-type]
-        elif isinstance(obj, staticmethod):
-            return staticmethod(decorator(obj.__func__))
-
-        @functools.wraps(obj)
-        def cast_args(*args: Tuple, **kwargs: Dict) -> Any:
-            """
-            Add casting for query compiler arguments.
-
-            Parameters
-            ----------
-            *args : tuple
-                The function arguments.
-            **kwargs : dict
-                The function keyword arguments.
-
-            Returns
-            -------
-            Any
-            """
-            current_qc = args[0]
-            if isinstance(current_qc, BaseQueryCompiler):
-                kwargs = cast_nested_args_to_current_qc_type(kwargs, current_qc)
-                args = cast_nested_args_to_current_qc_type(args, current_qc)
-            return obj(*args, **kwargs)
-
-        return cast_args
-
-    return decorator
+    return cast_args

modin/tests/pandas/native_df_mode/test_binary.py

@@ -17,9 +17,11 @@
 import pytest
 
 from modin.config import NativeDataframeMode, NPartitions
-from modin.tests.pandas.native_df_mode.utils import eval_general_interop
+from modin.tests.pandas.native_df_mode.utils import (
+    create_test_df_in_defined_mode,
+    eval_general_interop,
+)
 from modin.tests.pandas.utils import (
-    create_test_dfs,
     default_to_pandas_ignore_string,
     df_equals,
     test_data,
@@ -153,8 +155,12 @@ def operation(df1, df2):
     "df_mode_pair", list(product(NativeDataframeMode.choices, repeat=2))
 )
 def test_equals(frame1_data, frame2_data, expected_pandas_equals, df_mode_pair):
-    modin_df1, pandas_df1 = create_test_dfs(frame1_data, df_mode=df_mode_pair[0])
-    modin_df2, pandas_df2 = create_test_dfs(frame2_data, df_mode=df_mode_pair[1])
+    modin_df1, pandas_df1 = create_test_df_in_defined_mode(
+        frame1_data, df_mode=df_mode_pair[0]
+    )
+    modin_df2, pandas_df2 = create_test_df_in_defined_mode(
+        frame2_data, df_mode=df_mode_pair[1]
+    )
 
     pandas_equals = pandas_df1.equals(pandas_df2)
     assert pandas_equals == expected_pandas_equals, (
@@ -172,8 +178,12 @@ def test_equals(frame1_data, frame2_data, expected_pandas_equals, df_mode_pair):
     "df_mode_pair", list(product(NativeDataframeMode.choices, repeat=2))
 )
 def test_empty_df(empty_operand, df_mode_pair):
-    modin_df, pandas_df = create_test_dfs([0, 1, 2, 0, 1, 2], df_mode=df_mode_pair[0])
-    modin_df_empty, pandas_df_empty = create_test_dfs(df_mode=df_mode_pair[1])
+    modin_df, pandas_df = create_test_df_in_defined_mode(
+        [0, 1, 2, 0, 1, 2], df_mode=df_mode_pair[0]
+    )
+    modin_df_empty, pandas_df_empty = create_test_df_in_defined_mode(
+        df_mode=df_mode_pair[1]
+    )
 
     if empty_operand == "right":
         modin_res = modin_df + modin_df_empty

modin/tests/pandas/native_df_mode/test_default.py

@@ -23,10 +23,12 @@
 import modin.pandas as pd
 from modin.config import NativeDataframeMode, NPartitions
 from modin.pandas.io import to_pandas
-from modin.tests.pandas.native_df_mode.utils import eval_general_interop
+from modin.tests.pandas.native_df_mode.utils import (
+    create_test_df_in_defined_mode,
+    create_test_series_in_defined_mode,
+    eval_general_interop,
+)
 from modin.tests.pandas.utils import (
-    create_test_dfs,
-    create_test_series,
     default_to_pandas_ignore_string,
     df_equals,
     test_data,
@@ -81,12 +83,12 @@
     "df_mode_pair", list(product(NativeDataframeMode.choices, repeat=2))
 )
 def test_ops_defaulting_to_pandas(op, make_args, df_mode_pair):
-    modin_df1, _ = create_test_dfs(
+    modin_df1, _ = create_test_df_in_defined_mode(
         test_data_diff_dtype,
         post_fn=lambda df: df.drop(["str_col", "bool_col"], axis=1),
         df_mode=df_mode_pair[0],
     )
-    modin_df2, _ = create_test_dfs(
+    modin_df2, _ = create_test_df_in_defined_mode(
         test_data_diff_dtype,
         post_fn=lambda df: df.drop(["str_col", "bool_col"], axis=1),
         df_mode=df_mode_pair[1],
@@ -118,10 +120,10 @@ def test_to_numpy(data):
 )
 def test_asfreq(df_mode_pair):
     index = pd.date_range("1/1/2000", periods=4, freq="min")
-    series, _ = create_test_series(
+    series, _ = create_test_series_in_defined_mode(
         [0.0, None, 2.0, 3.0], index=index, df_mode=df_mode_pair[0]
     )
-    df, _ = create_test_dfs({"s": series}, df_mode=df_mode_pair[1])
+    df, _ = create_test_df_in_defined_mode({"s": series}, df_mode=df_mode_pair[1])
     with warns_that_defaulting_to_pandas():
         # We are only testing that this defaults to pandas, so we will just check for
         # the warning
@@ -152,9 +154,13 @@ def assign_multiple_columns(df1, df2):
 )
 def test_combine_first(df_mode_pair):
     data1 = {"A": [None, 0], "B": [None, 4]}
-    modin_df1, pandas_df1 = create_test_dfs(data1, df_mode=df_mode_pair[0])
+    modin_df1, pandas_df1 = create_test_df_in_defined_mode(
+        data1, df_mode=df_mode_pair[0]
+    )
     data2 = {"A": [1, 1], "B": [3, 3]}
-    modin_df2, pandas_df2 = create_test_dfs(data2, df_mode=df_mode_pair[1])
+    modin_df2, pandas_df2 = create_test_df_in_defined_mode(
+        data2, df_mode=df_mode_pair[1]
+    )
 
     df_equals(
         modin_df1.combine_first(modin_df2),
@@ -170,11 +176,11 @@ def test_combine_first(df_mode_pair):
 )
 def test_dot(data, df_mode_pair):
 
-    modin_df, pandas_df = create_test_dfs(data, df_mode=df_mode_pair[0])
+    modin_df, pandas_df = create_test_df_in_defined_mode(data, df_mode=df_mode_pair[0])
     col_len = len(modin_df.columns)
 
     # Test series input
-    modin_series, pandas_series = create_test_series(
+    modin_series, pandas_series = create_test_series_in_defined_mode(
         np.arange(col_len),
         index=pandas_df.columns,
         df_mode=df_mode_pair[1],
@@ -194,7 +200,7 @@ def dot_func(df1, df2):
 
     # Test when input series index doesn't line up with columns
     with pytest.raises(ValueError):
-        modin_series_without_index, _ = create_test_series(
+        modin_series_without_index, _ = create_test_series_in_defined_mode(
             np.arange(col_len), df_mode=df_mode_pair[1]
         )
         modin_df.dot(modin_series_without_index)
@@ -209,7 +215,7 @@ def dot_func(df1, df2):
     "df_mode_pair", list(product(NativeDataframeMode.choices, repeat=2))
 )
 def test_matmul(data, df_mode_pair):
-    modin_df, pandas_df = create_test_dfs(data, df_mode=df_mode_pair[0])
+    modin_df, pandas_df = create_test_df_in_defined_mode(data, df_mode=df_mode_pair[0])
     col_len = len(modin_df.columns)
 
     # Test list input
@@ -223,7 +229,7 @@ def test_matmul(data, df_mode_pair):
         modin_df @ np.arange(col_len + 10)
 
     # Test series input
-    modin_series, pandas_series = create_test_series(
+    modin_series, pandas_series = create_test_series_in_defined_mode(
         np.arange(col_len),
         index=pandas_df.columns,
         df_mode=df_mode_pair[1],
@@ -241,7 +247,7 @@ def matmul_func(df1, df2):
 
     # Test when input series index doesn't line up with columns
     with pytest.raises(ValueError):
-        modin_series_without_index, _ = create_test_series(
+        modin_series_without_index, _ = create_test_series_in_defined_mode(
             np.arange(col_len), df_mode=df_mode_pair[1]
         )
         modin_df @ modin_series_without_index