Add ;p pretty print option

PiperOrigin-RevId: 582094338

CHANGELOG.md

@@ -11,7 +11,8 @@ Changelog follow https://keepachangelog.com/ format.
 *   `ecolab`:
     *   Added protobuf repeated fields support to `ecolab.inspect`
     *   `ecolab.auto_display`:
-        *   Support specifiers to customize auto-display (`;s`, `;a`, `;i`,...)
+        *   Support specifiers to customize auto-display (`;s`, `;a`, `;i`,
+            `;p`,...)
         *   Fixed auto-display when the line contain UTF-8 character
 
 ## [1.5.2] - 2023-10-24

etils/ecolab/README.md

@@ -80,6 +80,9 @@ Format:
 *   `my_obj;i`: (`inspect`) Alias for `ecolab.inspect(x)`
 *   `my_obj;a`: (`array`) Alias for `media.show_images(x)` /
     `media.show_videos(x)` (`ecolab.auto_plot_array` behavior)
+*   `my_obj;p`: (`pretty_display`) Alias for `print(epy.pretty_repr(x))`.
+    Can be combined with `s`. Used for pretty print `dataclasses` or print
+    strings containing new lines (rather than displaying `\n`).
 *   `my_obj;q`: (`quiet`) Don't display the line (e.g. last line)
 
 ```python

etils/ecolab/auto_display_utils.py

@@ -40,7 +40,7 @@
 
 
 def auto_display(activate: bool = True) -> None:
-  """Activate auto display statements ending with `;` (activated by default).
+  r"""Activate auto display statements ending with `;` (activated by default).
 
   Add a trailing `;` to any statement (assignment, expression, return
   statement) to display the current line.
@@ -65,8 +65,13 @@ def auto_display(activate: bool = True) -> None:
   *   `my_obj;i`: (`inspect`) Alias for `ecolab.inspect(x)`
   *   `my_obj;a`: (`array`) Alias for `media.show_images(x)` /
       `media.show_videos(x)` (`ecolab.auto_plot_array` behavior)
+  *   `my_obj;p`: (`pretty_display`) Alias for `print(epy.pretty_repr(x))`.
+      Can be combined with `s`. Used for pretty print `dataclasses` or print
+      strings containing new lines (rather than displaying `\n`).
   *   `my_obj;q`: (`quiet`) Don't display the line (e.g. last line)
 
+  `p` and `s` can be combined.
+
   Args:
     activate: Allow to disable `auto_display`
   """
@@ -287,8 +292,9 @@ def _detect_trailing_regex() -> re.Pattern[str]:
   # Do not match:
   # * `; a; b`
   # * `; a=1`
-  available_chars = ''.join(_ALIAS_TO_DISPLAY_FN)
-  return re.compile(f' *; *([{available_chars}])? *(?:#.*)?$')
+
+  available_suffixes = '|'.join(list(_ALIAS_TO_DISPLAY_FN)[1:])
+  return re.compile(f' *; *({available_suffixes})? *(?:#.*)?$')
 
 
 def _display_and_return(x: _T) -> _T:
@@ -319,6 +325,24 @@ def _display_array_and_return(x: _T) -> _T:
   return x
 
 
+def _pretty_display_return(x: _T) -> _T:
+  """Print `x` and return `x`."""
+  # 2 main use-case:
+  # * Print strings (including `\n`)
+  # * Pretty-print dataclasses
+  if isinstance(x, str):
+    print(x)
+  else:
+    print(epy.pretty_repr(x))
+  return x
+
+
+def _pretty_display_specs_and_return(x: _T) -> _T:
+  """Print `x` and return `x`."""
+  print(epy.pretty_repr(etree.spec_like(x)))
+  return x
+
+
 def _return_quietly(x: _T) -> _T:
   """Return `x` without display."""
   return x
@@ -329,5 +353,8 @@ def _return_quietly(x: _T) -> _T:
     's': _display_specs_and_return,
     'i': _inspect_and_return,
     'a': _display_array_and_return,
+    'p': _pretty_display_return,
+    'ps': _pretty_display_specs_and_return,
+    'sp': _pretty_display_specs_and_return,
     'q': _return_quietly,
 }

etils/ecolab/auto_display_utils_test.py

@@ -93,6 +93,8 @@ def test_parsing():
         ('x;', ''),
         ('x;#i', ''),
         ('x;a', 'a'),
+        ('x;sp', 'sp'),
+        ('x;ps', 'ps'),
         ('x;i', 'i'),
         ('x   ;i', 'i'),
         ('x   ;  i', 'i'),
@@ -116,6 +118,8 @@ def test_alias(code: str, alias: str):
     'code',
     [
         'x;aa',
+        'x;ap',
+        'x;spp',
         'x;a=1',
     ],
 )