From e954c9a79f1d4ef4cd15f2364e5724f96e35cfff Mon Sep 17 00:00:00 2001
From: Richard Iannone <riannone@me.com>
Date: Tue, 24 Oct 2023 13:02:11 -0400
Subject: [PATCH] Add `_get_column_of_values()` util fn

---
 gt/gt.py              | 16 +++++++++++++++-
 tests/test_formats.py | 23 ++++++++++++++++-------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/gt/gt.py b/gt/gt.py
index 5f7b675fb..29d613575 100644
--- a/gt/gt.py
+++ b/gt/gt.py
@@ -225,10 +225,24 @@ def _set_has_built(gt: GT, value: bool) -> GT:
 
 def _get_column_labels(gt: GT, context: str) -> List[str]:
     gt_built = gt._build_data(context=context)
-    column_labels = gt_built._boxhead._get_column_labels()
+    # column_labels = gt_built._boxhead._get_column_labels()
+    column_labels = [x.column_label for x in gt_built._boxhead._boxhead]
     return column_labels
 
 
+def _get_column_of_values(gt: GT, column_name: str, context: str) -> List[str]:
+    gt_built = gt._build_data(context=context)
+    tbl_data = gt_built._body.body
+    cell_values: List[str] = []
+
+    for i in range(n_rows(tbl_data)):
+        cell_content: Any = _get_cell(tbl_data, i, column_name)
+        cell_str: str = str(cell_content)
+        cell_values.append(cell_str)
+
+    return cell_values
+
+
 # =============================================================================
 # Table Structuring Functions
 # =============================================================================
diff --git a/tests/test_formats.py b/tests/test_formats.py
index 49a72fd35..455d31f58 100644
--- a/tests/test_formats.py
+++ b/tests/test_formats.py
@@ -1,17 +1,26 @@
 import pandas as pd
-import pytest
-from pandas.testing import assert_frame_equal
 
 from gt import GT
+from gt.gt import _get_column_of_values
 
 
-@pytest.mark.xfail(reason="Don't have a method to get formatted values as list")
 def test_fmt_number_basic():
     df = pd.DataFrame({"x": [1.234, 2.345], "y": [3.456, 4.567]})
+
+    # Expect that values in `x` are formatted to 2 decimal places
     gt = GT(df).fmt_number(columns="x", decimals=2)
+    x = _get_column_of_values(gt, column_name="x", context="html")
+    y = ["1.23", "2.35"]
+    assert x == y
 
-    # TODO: is 2.35 below the intended result?
-    res = gt._build_data("html")._body.body
-    dst = df.assign(x=["1.23", "2.35"])
+    # TODO: this fails because unformatted values not migrated to body
+    # Expect that values in `y` are formatted to 2 decimal places
+    # x = _get_column_of_values(gt, column_name="y", context="html")
+    # y = ["3.46", "4.57"] # is currently ['<NA>', '<NA>']
+    # assert x == y
 
-    assert_frame_equal(res, dst)
+    # Expect that values in `x` are formatted to 5 decimal places
+    gt = GT(df).fmt_number(columns="x", decimals=5)
+    x = _get_column_of_values(gt, column_name="x", context="html")
+    y = ["1.23400", "2.34500"]
+    assert x == y