@@ -253,29 +253,57 @@ def test_build_df_using_interchange_protocol_mock(
253
253
add_interchange_module_for_old_pandas ,
254
254
):
255
255
class InterchangeDataFrame :
256
- def column_names (self ):
257
- return []
256
+ def __init__ (self , columns ):
257
+ self . _columns = columns
258
258
259
- def select_columns_by_name (self , column_names ):
260
- return self
259
+ def column_names (self ):
260
+ return self . _columns
261
261
262
- interchange_dataframe = InterchangeDataFrame ()
262
+ interchange_dataframe = InterchangeDataFrame (
263
+ ["petal_width" , "sepal_length" , "sepal_width" ]
264
+ )
265
+ interchange_dataframe_reduced = InterchangeDataFrame (
266
+ ["petal_width" , "sepal_length" ]
267
+ )
268
+ interchange_dataframe .select_columns_by_name = mock .MagicMock (
269
+ return_value = interchange_dataframe_reduced
270
+ )
271
+ interchange_dataframe_reduced .select_columns_by_name = mock .MagicMock (
272
+ return_value = interchange_dataframe_reduced
273
+ )
263
274
264
275
class CustomDataFrame :
265
276
def __dataframe__ (self ):
266
277
return interchange_dataframe
267
278
279
+ class CustomDataFrameReduced :
280
+ def __dataframe__ (self ):
281
+ return interchange_dataframe_reduced
282
+
268
283
input_dataframe = CustomDataFrame ()
269
- args = dict ( data_frame = input_dataframe , x = "petal_width" , y = "sepal_length" )
284
+ input_dataframe_reduced = CustomDataFrameReduced ( )
270
285
271
286
iris_pandas = px .data .iris ()
272
287
273
288
with mock .patch ("pandas.__version__" , "2.0.2" ):
289
+ args = dict (data_frame = input_dataframe , x = "petal_width" , y = "sepal_length" )
274
290
with mock .patch (
275
291
"pandas.api.interchange.from_dataframe" , return_value = iris_pandas
276
292
) as mock_from_dataframe :
277
293
build_dataframe (args , go .Scatter )
278
- mock_from_dataframe .assert_called_once_with (interchange_dataframe )
294
+ mock_from_dataframe .assert_called_once_with (interchange_dataframe_reduced )
295
+ interchange_dataframe .select_columns_by_name .assert_called_with (
296
+ ["petal_width" , "sepal_length" ]
297
+ )
298
+
299
+ args = dict (data_frame = input_dataframe_reduced , color = None )
300
+ with mock .patch (
301
+ "pandas.api.interchange.from_dataframe" ,
302
+ return_value = iris_pandas [["petal_width" , "sepal_length" ]],
303
+ ) as mock_from_dataframe :
304
+ build_dataframe (args , go .Scatter )
305
+ mock_from_dataframe .assert_called_once_with (interchange_dataframe_reduced )
306
+ interchange_dataframe_reduced .select_columns_by_name .assert_not_called ()
279
307
280
308
281
309
@pytest .mark .skipif (
0 commit comments