add separate index/columns sparsify mechanics

pandas-dev · jreback · May 12, 2021 · May 9, 2021 · May 11, 2021 · May 11, 2021
commit d29cac1c01fcc1c288f73011045757eb19fc41c2
diff --git a/pandas/io/formats/style_render.py b/pandas/io/formats/style_render.py
@@ -132,11 +132,16 @@ def _compute(self):
  r = func(self)(*args, **kwargs)
  return r
 
- def _translate(self):
+ def _translate(self, sparsify_index=None, sparsify_cols=None):
  """
  Convert the DataFrame in `self.data` and the attrs from `_build_styles`
  into a dictionary of {head, body, uuid, cellstyle}.
  """
+ if sparsify_index is None:
+ sparsify_index = get_option("display.multi_sparse")
+ if sparsify_cols is None:
+ sparsify_cols = get_option("display.multi_sparse")
+
  ROW_HEADING_CLASS = "row_heading"
  COL_HEADING_CLASS = "col_heading"
  INDEX_NAME_CLASS = "index_name"
@@ -153,14 +158,14 @@ def _translate(self):
  }
 
  head = self._translate_header(
- BLANK_CLASS, BLANK_VALUE, INDEX_NAME_CLASS, COL_HEADING_CLASS
+ BLANK_CLASS, BLANK_VALUE, INDEX_NAME_CLASS, COL_HEADING_CLASS, sparsify_cols
  )
  d.update({"head": head})
 
  self.cellstyle_map: DefaultDict[tuple[CSSPair, ...], list[str]] = defaultdict(
  list
  )
- body = self._translate_body(DATA_CLASS, ROW_HEADING_CLASS)
+ body = self._translate_body(DATA_CLASS, ROW_HEADING_CLASS, sparsify_index)
  d.update({"body": body})
 
  cellstyle: list[dict[str, CSSList | list[str]]] = [
@@ -185,7 +190,12 @@ def _translate(self):
  return d
 
  def _translate_header(
- self, blank_class, blank_value, index_name_class, col_heading_class
+ self,
+ blank_class,
+ blank_value,
+ index_name_class,
+ col_heading_class,
+ sparsify_cols,
  ):
  """
  Build each <tr> within table <head>, using the structure:
@@ -198,7 +208,9 @@ def _translate_header(
  +----------------------------+---------------+---------------------------+
  """
  # for sparsifying a MultiIndex
- col_lengths = _get_level_lengths(self.columns, self.hidden_columns)
+ col_lengths = _get_level_lengths(
+ self.columns, sparsify_cols, self.hidden_columns
+ )
 
  clabels = self.data.columns.tolist()
  if self.data.columns.nlevels == 1:
@@ -268,7 +280,7 @@ def _translate_header(
 
  return head
 
- def _translate_body(self, data_class, row_heading_class):
+ def _translate_body(self, data_class, row_heading_class, sparsify_index):
  """
  Build each <tr> in table <body> in the following format:
  +--------------------------------------------+---------------------------+
@@ -279,7 +291,7 @@ def _translate_body(self, data_class, row_heading_class):
  <style></style> block
  """
  # for sparsifying a MultiIndex
- idx_lengths = _get_level_lengths(self.index)
+ idx_lengths = _get_level_lengths(self.index, sparsify_index)
 
  rlabels = self.data.index.tolist()
  if self.data.index.nlevels == 1:
@@ -520,14 +532,23 @@ def _element(
  }
 
 
-def _get_level_lengths(index, hidden_elements=None):
+def _get_level_lengths(index, sparsify, hidden_elements=None):
  """
  Given an index, find the level length for each element.
 
- Optional argument is a list of index positions which
- should not be visible.
-
- Result is a dictionary of (level, initial_position): span
+ Parameters
+ ----------
+ index : Index
+ Index or columns to determine lengths of each element
+ sparsify : bool
+ Whether to hide or show each distinct element in a MultiIndex
+ hidden_element : list
+ Index positions of elements hidden from display in the index affecting
+ length
+ Returns
+ -------
+ Dict :
+ Result is a dictionary of (level, initial_position): span
  """
  if isinstance(index, MultiIndex):
  levels = index.format(sparsify=lib.no_default, adjoin=False)
@@ -546,7 +567,7 @@ def _get_level_lengths(index, hidden_elements=None):
 
  for i, lvl in enumerate(levels):
  for j, row in enumerate(lvl):
- if not get_option("display.multi_sparse"):
+ if not sparsify:
  lengths[(i, j)] = 1
  elif (row is not lib.no_default) and (j not in hidden_elements):
  last_label = j

diff --git a/pandas/tests/io/formats/style/test_style.py b/pandas/tests/io/formats/style/test_style.py
@@ -844,7 +844,24 @@ def test_get_level_lengths(self):
  (1, 4): 1,
  (1, 5): 1,
  }
- result = _get_level_lengths(index)
+ result = _get_level_lengths(index, sparsify=True)
+ tm.assert_dict_equal(result, expected)
+
+ expected = {
+ (0, 0): 1,
+ (0, 1): 1,
+ (0, 2): 1,
+ (0, 3): 1,
+ (0, 4): 1,
+ (0, 5): 1,
+ (1, 0): 1,
+ (1, 1): 1,
+ (1, 2): 1,
+ (1, 3): 1,
+ (1, 4): 1,
+ (1, 5): 1,
+ }
+ result = _get_level_lengths(index, sparsify=False)
  tm.assert_dict_equal(result, expected)
 
  def test_get_level_lengths_un_sorted(self):
@@ -858,7 +875,20 @@ def test_get_level_lengths_un_sorted(self):
  (1, 2): 1,
  (1, 3): 1,
  }
- result = _get_level_lengths(index)
+ result = _get_level_lengths(index, sparsify=True)
+ tm.assert_dict_equal(result, expected)
+
+ expected = {
+ (0, 0): 1,
+ (0, 1): 1,
+ (0, 2): 1,
+ (0, 3): 1,
+ (1, 0): 1,
+ (1, 1): 1,
+ (1, 2): 1,
+ (1, 3): 1,
+ }
+ result = _get_level_lengths(index, sparsify=False)
  tm.assert_dict_equal(result, expected)
 
  def test_mi_sparse(self):