From 2dc633226fe4f24dd2f8c56d2190925c6e9e94a9 Mon Sep 17 00:00:00 2001
From: Michael Chow <mc_al_github@fastmail.com>
Date: Mon, 12 Aug 2024 14:30:15 -0700
Subject: [PATCH] fix: include subset_frame for gt render

---
 reactable/_tbl_data.py | 32 +++++++++++++++++++++++++++++++-
 reactable/render_gt.py |  3 ++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/reactable/_tbl_data.py b/reactable/_tbl_data.py
index 6d01141..d214123 100644
--- a/reactable/_tbl_data.py
+++ b/reactable/_tbl_data.py
@@ -4,7 +4,7 @@
 
 from datetime import datetime, date, time
 from functools import singledispatch
-from typing import TYPE_CHECKING, Any, Union, Literal
+from typing import TYPE_CHECKING, Any, Union, Literal, Optional
 from typing_extensions import TypeAlias
 
 from abc import ABC
@@ -163,3 +163,33 @@ def _(data: PdDataFrame) -> "list[str]":
 @column_names.register
 def _(data: SimpleFrame) -> "list[str]":
     return list(data.columns)
+
+
+# subset_frame --------------------------------------------------------
+@singledispatch
+def subset_frame(
+    data: DataFrameLike, row: Optional[list[int]], column: Optional[list[str]]
+) -> DataFrameLike:
+    raise TypeError(f"Unsupported type: {type(data)}")
+
+
+@subset_frame.register
+def _(
+    data: PdDataFrame, rows: Optional[list[int]] = None, cols: Optional[list[str]] = None
+) -> PdDataFrame:
+
+    cols_indx = slice(None) if cols is None else data.columns.get_indexer_for(cols)
+    rows_indx = slice(None) if rows is None else rows
+
+    return data.iloc[rows_indx, cols_indx]
+
+
+@subset_frame.register
+def _(
+    data: PlDataFrame, rows: Optional[list[int]] = None, cols: Optional[list[str]] = None
+) -> PlDataFrame:
+
+    cols_indx = slice(None) if cols is None else cols
+    rows_indx = slice(None) if rows is None else rows
+
+    return data[rows_indx, cols_indx]
diff --git a/reactable/render_gt.py b/reactable/render_gt.py
index 2dc3274..2f1bdce 100644
--- a/reactable/render_gt.py
+++ b/reactable/render_gt.py
@@ -5,7 +5,7 @@
 import ipywidgets
 
 from great_tables import GT
-from great_tables._tbl_data import n_rows, subset_frame
+from great_tables._tbl_data import n_rows
 from great_tables._helpers import random_id
 from great_tables._text import _process_text
 from great_tables._gt_data import ColInfoTypeEnum
@@ -15,6 +15,7 @@
 from .models import Column, Language, Theme, ColGroup
 from . import Reactable
 from .tags import as_react_style, to_widget
+from ._tbl_data import subset_frame
 
 if TYPE_CHECKING:
     from great_tables._gt_data import Locale, Spanners, Heading, Footnotes, SourceNotes