Merge pull request #4442 from anmyachev/fix-column-names

Fix usage of `column_names` which is part of Interchange protocol

Author: alexcjohnson

CHANGELOG.md

@@ -9,6 +9,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 - Fix issue with creating dendrogram in subplots [[#4411](https://github.com/plotly/plotly.py/pull/4411)],
 - Fix issue with px.line not accepting "spline" line shape [[#2812](https://github.com/plotly/plotly.py/issues/2812)]
 - Fix KeyError when using column of `pd.Categorical` dtype with unobserved categories [[#4437](https://github.com/plotly/plotly.py/pull/4437)]
+- Fix dataframe interchange in case `column_names` returns an unmaterialized object: generator, iterator, etc. [[#4442]](https://github.com/plotly/plotly.py/pull/4442)
 
 ## [5.18.0] - 2023-10-25
 

packages/python/plotly/plotly/express/_core.py

@@ -1327,7 +1327,11 @@ def build_dataframe(args, constructor):
 
             df_not_pandas = args["data_frame"]
             args["data_frame"] = df_not_pandas.__dataframe__()
-            columns = args["data_frame"].column_names()
+            # According interchange protocol: `def column_names(self) -> Iterable[str]:`
+            # so this function can return for example a generator.
+            # The easiest way is to convert `columns` to `pandas.Index` so that the
+            # type is similar to the types in other code branches.
+            columns = pd.Index(args["data_frame"].column_names())
             needs_interchanging = True
         elif hasattr(args["data_frame"], "to_pandas"):
             args["data_frame"] = args["data_frame"].to_pandas()

packages/python/plotly/plotly/tests/test_optional/test_px/test_px_input.py

@@ -14,12 +14,12 @@
 @pytest.fixture
 def add_interchange_module_for_old_pandas():
     if not hasattr(pd.api, "interchange"):
-        pd.api.interchange = mock.MagicMock()
-        # to make the following import work: `import pandas.api.interchange`
-        with mock.patch.dict(
-            "sys.modules", {"pandas.api.interchange": pd.api.interchange}
-        ):
-            yield
+        with mock.patch.object(pd.api, "interchange", mock.MagicMock(), create=True):
+            # to make the following import work: `import pandas.api.interchange`
+            with mock.patch.dict(
+                "sys.modules", {"pandas.api.interchange": pd.api.interchange}
+            ):
+                yield
     else:
         yield
 
@@ -250,15 +250,24 @@ def test_build_df_with_index():
     assert_frame_equal(tips.reset_index()[out["data_frame"].columns], out["data_frame"])
 
 
+@pytest.mark.parametrize("column_names_as_generator", [False, True])
 def test_build_df_using_interchange_protocol_mock(
-    add_interchange_module_for_old_pandas,
+    add_interchange_module_for_old_pandas, column_names_as_generator
 ):
     class InterchangeDataFrame:
         def __init__(self, columns):
             self._columns = columns
 
-        def column_names(self):
-            return self._columns
+        if column_names_as_generator:
+
+            def column_names(self):
+                for col in self._columns:
+                    yield col
+
+        else:
+
+            def column_names(self):
+                return self._columns
 
     interchange_dataframe = InterchangeDataFrame(
         ["petal_width", "sepal_length", "sepal_width"]