Fix slicing capabilities (#206)

* fix slicing capabilities

* [pre-commit.ci lite] apply automatic fixes

* linting

* Update .pre-commit-config.yaml

* address reviews

* address coderabbit review

* linting

* Update .pre-commit-config.yaml

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: d33bs

.pre-commit-config.yaml

@@ -50,7 +50,7 @@ repos:
     hooks:
     -   id: actionlint
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.15.10"
+    rev: "v0.15.12"
     hooks:
     -   id: ruff-format
     -   id: ruff-check

src/cytodataframe/frame.py

@@ -124,7 +124,7 @@ def __init__(  # noqa: PLR0913
         segmentation_file_regex: Optional[Dict[str, str]] = None,
         image_adjustment: Optional[Callable] = None,
         display_options: Optional[Dict[str, Any]] = None,
-        *args: Tuple[Any, ...],
+        *args: Any,
         **kwargs: Dict[str, Any],
     ) -> None:
         """
@@ -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,34 +373,56 @@ 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: 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,
+            *args,
+            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"],
+            **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,
         method: Callable,
         method_name: str,
-        *args: Tuple[Any, ...],
+        *args: Any,
         **kwargs: Dict[str, Any],
     ) -> Any:
         """
@@ -410,7 +434,7 @@ def _return_cytodataframe(
                 The method to be called and wrapped.
             method_name (str):
                 The name of the method to be wrapped.
-            *args (Tuple[Any, ...]):
+            *args (Any):
                 Positional arguments to be passed to the method.
             **kwargs (Dict[str, Any]):
                 Keyword arguments to be passed to the method.
@@ -423,33 +447,20 @@ def _return_cytodataframe(
 
         """
 
-        result = 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"],
-            )
-            # 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
+        result = (
+            pd.DataFrame(self).transpose(*args, **kwargs)
+            if method_name == "transpose"
+            else method(*args, **kwargs)
+        )
 
-            # 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"]
+        if not isinstance(result, pd.DataFrame):
+            return result
+
+        cdf = self._build_result_cdf(result)
+        # If the method name is transpose we know that
+        # the dataframe has been transposed.
+        if method_name == "transpose":
+            cdf._custom_attrs["is_transposed"] = not self._custom_attrs["is_transposed"]
 
         return cdf
 
@@ -473,7 +484,7 @@ def _wrap_method(self: CytoDataFrame_type, method_name: str) -> Callable:
                 the result is a CytoDataFrame.
         """
 
-        def wrapper(*args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any:
+        def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
             """
             Wraps the specified method to ensure
             it returns a CytoDataFrame.
@@ -485,7 +496,7 @@ def wrapper(*args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any:
             custom attributes.
 
             Args:
-                *args (Tuple[Any, ...]):
+                *args (Any):
                     Positional arguments to be passed to the method.
                 **kwargs (Dict[str, Any]):
                     Keyword arguments to be passed to the method.
@@ -4346,7 +4357,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 +4373,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 +4397,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 +4428,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,71 @@ 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[0:2], CytoDataFrame)
+    assert isinstance(cdf[1:1], CytoDataFrame)
+    assert isinstance(cdf[0:5:2], 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_return_cytodataframe_passthroughs_non_dataframe_results() -> None:
+    """Ensure helper methods return scalar-like results without wrapping."""
+
+    cdf = CytoDataFrame(pd.DataFrame({"a": [1, 2, 3]}))
+
+    result = cdf._return_cytodataframe(lambda: 3, "dummy_method")
+
+    assert result == 3
+
+
+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]}))
+
+    bracket_sliced = cdf[0:3:2]
+    bracket_empty_sliced = cdf[1:1]
+    sliced = cdf.iloc[0:3:2]
+    empty_sliced = cdf.iloc[1:1]
+
+    assert isinstance(bracket_sliced, CytoDataFrame)
+    assert isinstance(bracket_empty_sliced, CytoDataFrame)
+    assert bracket_sliced._custom_attrs["_output"] is cdf._custom_attrs["_output"]
+    assert (
+        bracket_sliced._custom_attrs["_widget_state"]
+        is cdf._custom_attrs["_widget_state"]
+    )
+    assert bracket_empty_sliced._custom_attrs["_output"] is cdf._custom_attrs["_output"]
+    assert (
+        bracket_empty_sliced._custom_attrs["_widget_state"]
+        is cdf._custom_attrs["_widget_state"]
+    )
+    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 bracket_sliced._repr_html_(debug=True)
+    assert "background:#EBEBEB" in bracket_empty_sliced._repr_html_(debug=True)
+    assert "background:#EBEBEB" in sliced._repr_html_(debug=True)
+    assert "background:#EBEBEB" in empty_sliced._repr_html_(debug=True)
+
+
+def test_transpose_toggles_transposed_state() -> None:
+    """Ensure repeated transposes flip the transposed rendering state back."""
+
+    cdf = CytoDataFrame(pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
+
+    transposed = cdf.T
+    double_transposed = transposed.T
+
+    assert transposed._custom_attrs["is_transposed"] is True
+    assert double_transposed._custom_attrs["is_transposed"] is False
 
 
 def test_cytodataframe_dynamic_width_and_height(