fix slicing capabilities

Author: d33bs

src/cytodataframe/frame.py

@@ -277,6 +277,7 @@ def __init__(  # noqa: PLR0913
                 )
             ),
             "_filter_range_sliders": {},
+            "_initializing": True,
         }
 
         if self._custom_attrs["data_context_dir"] is not None:
@@ -351,6 +352,7 @@ def __init__(  # noqa: PLR0913
         # Wrap methods so they return CytoDataFrames
         # instead of Pandas DataFrames.
         self._wrap_methods()
+        self._custom_attrs["_initializing"] = False
 
     def __getitem__(self: CytoDataFrame_type, key: Union[int, str]) -> Any:
         """
@@ -371,28 +373,50 @@ def __getitem__(self: CytoDataFrame_type, key: Union[int, str]) -> Any:
             return result
 
         elif isinstance(result, pd.DataFrame):
-            cdf = CytoDataFrame(
-                super().__getitem__(key),
-                data_context_dir=self._custom_attrs["data_context_dir"],
-                data_image_paths=self._custom_attrs["data_image_paths"],
-                data_bounding_box=self._custom_attrs["data_bounding_box"],
-                compartment_center_xy=self._custom_attrs["compartment_center_xy"],
-                data_mask_context_dir=self._custom_attrs["data_mask_context_dir"],
-                data_outline_context_dir=self._custom_attrs["data_outline_context_dir"],
-                segmentation_file_regex=self._custom_attrs["segmentation_file_regex"],
-                image_adjustment=self._custom_attrs["image_adjustment"],
-                display_options=self._custom_attrs["display_options"],
-            )
+            return self._build_result_cdf(result)
 
-            # add widget control meta
-            cdf._custom_attrs["_widget_state"] = self._custom_attrs["_widget_state"]
-            cdf._custom_attrs["_scale_slider"] = self._custom_attrs["_scale_slider"]
-            cdf._custom_attrs["_filter_range_sliders"] = self._custom_attrs[
-                "_filter_range_sliders"
-            ]
-            cdf._custom_attrs["_output"] = self._custom_attrs["_output"]
+    @property
+    def _constructor(self: CytoDataFrame_type) -> Callable[..., CytoDataFrame_type]:
+        """Return a constructor that preserves CytoDataFrame display state."""
+        if self._custom_attrs.get("_initializing", False):
+            return pd.DataFrame
+        return self._build_result_cdf
 
-            return cdf
+    def _build_result_cdf(
+        self: CytoDataFrame_type,
+        data: Union["CytoDataFrame", pd.DataFrame, pd.Series, Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> CytoDataFrame_type:
+        """Construct a result frame while preserving CytoDataFrame metadata."""
+        if data is None and "data" in kwargs:
+            data = kwargs.pop("data")
+
+        cdf = CytoDataFrame(
+            data=data,
+            data_context_dir=self._custom_attrs["data_context_dir"],
+            data_image_paths=self._custom_attrs["data_image_paths"],
+            data_bounding_box=self._custom_attrs["data_bounding_box"],
+            compartment_center_xy=self._custom_attrs["compartment_center_xy"],
+            data_mask_context_dir=self._custom_attrs["data_mask_context_dir"],
+            data_outline_context_dir=self._custom_attrs["data_outline_context_dir"],
+            segmentation_file_regex=self._custom_attrs["segmentation_file_regex"],
+            image_adjustment=self._custom_attrs["image_adjustment"],
+            display_options=self._custom_attrs["display_options"],
+            *args,
+            **kwargs,
+        )
+
+        # Preserve the shared interactive widget objects across derived views.
+        cdf._custom_attrs["is_transposed"] = self._custom_attrs["is_transposed"]
+        cdf._custom_attrs["_widget_state"] = self._custom_attrs["_widget_state"]
+        cdf._custom_attrs["_scale_slider"] = self._custom_attrs["_scale_slider"]
+        cdf._custom_attrs["_filter_range_sliders"] = self._custom_attrs[
+            "_filter_range_sliders"
+        ]
+        cdf._custom_attrs["_output"] = self._custom_attrs["_output"]
+
+        return cdf
 
     def _return_cytodataframe(
         self: CytoDataFrame_type,
@@ -423,34 +447,19 @@ def _return_cytodataframe(
 
         """
 
-        result = method(*args, **kwargs)
+        result = (
+            pd.DataFrame(self).transpose(*args, **kwargs)
+            if method_name == "transpose"
+            else method(*args, **kwargs)
+        )
 
         if isinstance(result, pd.DataFrame):
-            cdf = CytoDataFrame(
-                data=result,
-                data_context_dir=self._custom_attrs["data_context_dir"],
-                data_image_paths=self._custom_attrs["data_image_paths"],
-                data_bounding_box=self._custom_attrs["data_bounding_box"],
-                compartment_center_xy=self._custom_attrs["compartment_center_xy"],
-                data_mask_context_dir=self._custom_attrs["data_mask_context_dir"],
-                data_outline_context_dir=self._custom_attrs["data_outline_context_dir"],
-                segmentation_file_regex=self._custom_attrs["segmentation_file_regex"],
-                image_adjustment=self._custom_attrs["image_adjustment"],
-                display_options=self._custom_attrs["display_options"],
-            )
+            cdf = self._build_result_cdf(result)
             # If the method name is transpose we know that
             # the dataframe has been transposed.
             if method_name == "transpose" and not self._custom_attrs["is_transposed"]:
                 cdf._custom_attrs["is_transposed"] = True
 
-            # add widget control meta
-            cdf._custom_attrs["_widget_state"] = self._custom_attrs["_widget_state"]
-            cdf._custom_attrs["_scale_slider"] = self._custom_attrs["_scale_slider"]
-            cdf._custom_attrs["_filter_range_sliders"] = self._custom_attrs[
-                "_filter_range_sliders"
-            ]
-            cdf._custom_attrs["_output"] = self._custom_attrs["_output"]
-
         return cdf
 
     def _wrap_method(self: CytoDataFrame_type, method_name: str) -> Callable:
@@ -4346,7 +4355,7 @@ def _generate_jupyter_dataframe_html(  # noqa: C901, PLR0912, PLR0915
                 # if the data are transposed,
                 # we transpose them back to keep
                 # logic the same here.
-                data = self.transpose()
+                data = pd.DataFrame(self).transpose()
 
             # Re-add bounding box columns if they are no longer available
             bounding_box_externally_joined = False
@@ -4362,7 +4371,12 @@ def _generate_jupyter_dataframe_html(  # noqa: C901, PLR0912, PLR0915
                 )
                 bounding_box_externally_joined = True
             else:
-                data = self.copy() if not bounding_box_externally_joined else data
+                data = (
+                    data
+                    if self._custom_attrs["is_transposed"]
+                    or bounding_box_externally_joined
+                    else self.copy()
+                )
 
             # Re-add compartment center xy columns if they are no longer available
             compartment_center_externally_joined = False
@@ -4381,7 +4395,8 @@ def _generate_jupyter_dataframe_html(  # noqa: C901, PLR0912, PLR0915
             else:
                 data = (
                     data
-                    if bounding_box_externally_joined
+                    if self._custom_attrs["is_transposed"]
+                    or bounding_box_externally_joined
                     or compartment_center_externally_joined
                     else self.copy()
                 )
@@ -4411,7 +4426,9 @@ def _generate_jupyter_dataframe_html(  # noqa: C901, PLR0912, PLR0915
             else:
                 data = (
                     data
-                    if image_paths_externally_joined or bounding_box_externally_joined
+                    if self._custom_attrs["is_transposed"]
+                    or image_paths_externally_joined
+                    or bounding_box_externally_joined
                     else self.copy()
                 )
 

tests/test_frame.py

@@ -796,6 +796,30 @@ def test_return_cytodataframe(cytotable_NF1_data_parquet_shrunken: str):
     assert isinstance(cdf.tail(), CytoDataFrame)
     assert isinstance(cdf.sort_values(by="Metadata_ImageNumber"), CytoDataFrame)
     assert isinstance(cdf.sample(n=5), CytoDataFrame)
+    assert isinstance(cdf.iloc[0:2], CytoDataFrame)
+    assert isinstance(cdf.iloc[1:1], CytoDataFrame)
+    assert isinstance(cdf.iloc[0:5:2], CytoDataFrame)
+
+
+def test_iloc_slice_preserves_cytodataframe_html_formatting():
+    """Ensure ``iloc`` slices keep the CytoDataFrame notebook HTML renderer."""
+
+    cdf = CytoDataFrame(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+
+    sliced = cdf.iloc[0:3:2]
+    empty_sliced = cdf.iloc[1:1]
+
+    assert isinstance(sliced, CytoDataFrame)
+    assert isinstance(empty_sliced, CytoDataFrame)
+    assert sliced._custom_attrs["_output"] is cdf._custom_attrs["_output"]
+    assert sliced._custom_attrs["_widget_state"] is cdf._custom_attrs["_widget_state"]
+    assert empty_sliced._custom_attrs["_output"] is cdf._custom_attrs["_output"]
+    assert (
+        empty_sliced._custom_attrs["_widget_state"]
+        is cdf._custom_attrs["_widget_state"]
+    )
+    assert "background:#EBEBEB" in sliced._repr_html_(debug=True)
+    assert "background:#EBEBEB" in empty_sliced._repr_html_(debug=True)
 
 
 def test_cytodataframe_dynamic_width_and_height(