add tests, enhance doc string and formatter wrapping

pandas/io/formats/style.py

@@ -71,7 +71,7 @@ class Styler:
  The ``id`` takes the form ``T_<uuid>_row<num_row>_col<num_col>``
  where ``<uuid>`` is the unique identifier, ``<num_row>`` is the row
  number and ``<num_col>`` is the column number.
- na_rep : str or None, default None
+ na_rep : str, optional
  Representation for missing values.
  If ``na_rep`` is None, no special formatting is applied
 
@@ -436,7 +436,7 @@ def format(self, formatter, subset=None, na_rep=None):
  subset : IndexSlice
  An argument to ``DataFrame.loc`` that restricts which elements
  ``formatter`` is applied to.
- na_rep : str or None, default None
+ na_rep : str, optional
  Representation for missing values.
  If ``na_rep`` is None, no special formatting is applied
 
@@ -484,16 +484,14 @@ def format(self, formatter, subset=None, na_rep=None):
  if is_dict_like(formatter):
  for col, col_formatter in formatter.items():
  # formatter must be callable, so '{}' are converted to lambdas
- col_formatter = _maybe_wrap_formatter(col_formatter)
- col_formatter = _maybe_wrap_na_formatter(col_formatter, na_rep)
+ col_formatter = _maybe_wrap_formatter(col_formatter, na_rep)
  col_num = self.data.columns.get_indexer_for([col])[0]
 
  for row_num in row_locs:
  self._display_funcs[(row_num, col_num)] = col_formatter
  else:
  # single scalar to format all cells with
- formatter = _maybe_wrap_formatter(formatter)
- formatter = _maybe_wrap_na_formatter(formatter, na_rep)
+ formatter = _maybe_wrap_formatter(formatter, na_rep)
  locs = product(*(row_locs, col_locs))
  for i, j in locs:
  self._display_funcs[(i, j)] = formatter
@@ -1507,24 +1505,22 @@ def _get_level_lengths(index, hidden_elements=None):
  return non_zero_lengths
 
 
-def _maybe_wrap_formatter(formatter):
+def _maybe_wrap_formatter(formatter, na_rep):
  if is_string_like(formatter):
- return lambda x: formatter.format(x)
+ formatter_func = lambda x: formatter.format(x)
  elif callable(formatter):
- return formatter
+ formatter_func = formatter
  else:
  msg = (
  "Expected a template string or callable, got {formatter} "
  "instead".format(formatter=formatter)
  )
  raise TypeError(msg)
 
-
-def _maybe_wrap_na_formatter(formatter, na_rep):
  if na_rep is None:
- return formatter
+ return formatter_func
  elif is_string_like(na_rep):
- return lambda x: na_rep if pd.isna(x) else formatter(x)
+ return lambda x: na_rep if pd.isna(x) else formatter_func(x)
  else:
  msg = "Expected a string, got {na_rep} instead".format(na_rep=na_rep)
  raise TypeError(msg)

pandas/tests/io/formats/test_style.py

@@ -1008,6 +1008,14 @@ def test_format_with_na_rep(self):
  assert ctx["body"][0][2]["display_value"] == "-"
  assert ctx["body"][1][2]["display_value"] == "120.00%"
 
+ def test_init_with_na_rep(self):
+ # GH 21527 28358
+ df = pd.DataFrame([[None, None], [1.1, 1.2]], columns=["A", "B"])
+
+ ctx = Styler(df, na_rep="NA")._translate()
+ assert ctx["body"][0][1]["display_value"] == "NA"
+ assert ctx["body"][0][2]["display_value"] == "NA"
+
  def test_set_na_rep(self):
  # GH 21527 28358
  df = pd.DataFrame([[None, None], [1.1, 1.2]], columns=["A", "B"])
@@ -1024,6 +1032,27 @@ def test_set_na_rep(self):
  assert ctx["body"][0][1]["display_value"] == "NA"
  assert ctx["body"][0][2]["display_value"] == "-"
 
+ def test_format_non_numeric_na(self):
+ # GH 21527 28358
+ df = pd.DataFrame(
+ {
+ "object": [None, np.nan, "foo"],
+ "datetime": [None, pd.NaT, pd.Timestamp("20120101")],
+ }
+ )
+
+ ctx = df.style.set_na_rep("NA")._translate()
+ assert ctx["body"][0][1]["display_value"] == "NA"
+ assert ctx["body"][0][2]["display_value"] == "NA"
+ assert ctx["body"][1][1]["display_value"] == "NA"
+ assert ctx["body"][1][2]["display_value"] == "NA"
+
+ ctx = df.style.format(None, na_rep="-")._translate()
+ assert ctx["body"][0][1]["display_value"] == "-"
+ assert ctx["body"][0][2]["display_value"] == "-"
+ assert ctx["body"][1][1]["display_value"] == "-"
+ assert ctx["body"][1][2]["display_value"] == "-"
+
  def test_highlight_null(self, null_color="red"):
  df = pd.DataFrame({"A": [0, np.nan]})
  result = df.style.highlight_null()._compute().ctx