meerschaum.utils.formatting

  1#! /usr/bin/env python
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Utilities for formatting output text
  7"""
  8
  9from __future__ import annotations
 10import platform
 11import os
 12import sys
 13import meerschaum as mrsm
 14from meerschaum.utils.typing import Optional, Union, Any, Dict, Iterable
 15from meerschaum.utils.formatting._shell import make_header
 16from meerschaum.utils.formatting._pprint import pprint
 17from meerschaum.utils.formatting._dataframe import pprint_df, format_dataframe
 18from meerschaum.utils.formatting._pipes import (
 19    pprint_pipes,
 20    highlight_pipes,
 21    format_pipe_success_tuple,
 22    print_pipes_results,
 23    extract_stats_from_message,
 24    pipe_repr,
 25)
 26from meerschaum.utils.threading import Lock, RLock
 27
 28_attrs = {
 29    'ANSI': None,
 30    'UNICODE': None,
 31    'CHARSET': None,
 32    'RESET': '\033[0m',
 33}
 34__all__ = sorted([
 35    'ANSI', 'CHARSET', 'UNICODE', 'RESET',
 36    'ansi_allowed',
 37    'stderr_ansi',
 38    'colored',
 39    'translate_rich_to_termcolor',
 40    'get_console',
 41    'print_exception',
 42    'format_success_tuple',
 43    'print_tuple',
 44    'print_options',
 45    'fill_ansi',
 46    'pprint',
 47    'pprint_df',
 48    'format_dataframe',
 49    'highlight_pipes',
 50    'pprint_pipes',
 51    'make_header',
 52    'pipe_repr',
 53    'print_pipes_results',
 54    'extract_stats_from_message',
 55    'format_bytes',
 56])
 57__pdoc__ = {}
 58_locks = {
 59    '_colorama_init': RLock(),
 60}
 61
 62
 63def colored_fallback(*args, **kw):
 64    return ' '.join(args)
 65
 66
 67def _env_flag(name: str) -> str:
 68    """The trimmed, lower-case value of an environment variable (`''` when unset)."""
 69    return (os.environ.get(name, None) or '').strip().lower()
 70
 71
 72def ansi_allowed(stream: Optional[Any] = None) -> bool:
 73    """
 74    Return whether output may carry ANSI styles, before the `formatting:ansi` setting.
 75
 76    `NO_COLOR` (any non-empty value) turns styles off. Otherwise styles are on when the stream is
 77    a terminal, inside a job (its logs keep their colors, `MRSM_DAEMON_ID`), in a Jupyter notebook
 78    or when `FORCE_COLOR` is set (any value but `0`, `false` or `no`). Writing to a file or a
 79    pipe (`mrsm show pipes > pipes.txt`) therefore prints plain text.
 80
 81    Parameters
 82    ----------
 83    stream: Optional[Any], default None
 84        The stream the output is written to (`sys.stderr` for warnings). Defaults to `sys.stdout`.
 85    """
 86    if os.environ.get('NO_COLOR', None):
 87        return False
 88    force = _env_flag('FORCE_COLOR')
 89    if force and force not in ('0', 'false', 'no'):
 90        return True
 91    from meerschaum._internal.static import STATIC_CONFIG
 92    if os.environ.get(STATIC_CONFIG['environment']['daemon_id'], None):
 93        return True
 94    if _in_notebook():
 95        return True
 96    try:
 97        return bool((sys.stdout if stream is None else stream).isatty())
 98    except Exception:
 99        return False
100
101
102def stderr_ansi() -> bool:
103    """
104    Whether text written to stderr (warnings) may carry ANSI styles: `formatting:ansi` and
105    `ansi_allowed(sys.stderr)`, so `2> errors.log` gets plain text while stdout is a terminal.
106    """
107    __getattr__('ANSI')  ### Loads `formatting:ansi` into `_attrs`.
108    return bool(_attrs['ANSI']) and ansi_allowed(sys.stderr)
109
110
111def _in_notebook() -> bool:
112    """
113    Whether this is a Jupyter (IPython kernel) process: its outputs render ANSI styles although
114    its stdout is not a terminal. Only an already-imported `ipykernel` is looked at.
115    """
116    if 'ipykernel' not in sys.modules:
117        return False
118    try:
119        from IPython import get_ipython
120        shell = get_ipython()
121    except Exception:
122        return False
123    return shell is not None and getattr(shell, 'kernel', None) is not None
124
125
126def format_bytes(num_bytes: Optional[Union[int, float]], precision: int = 1) -> str:
127    """
128    Return a human-readable representation of a number of bytes.
129
130    Parameters
131    ----------
132    num_bytes: Optional[Union[int, float]]
133        The number of bytes to format. If `None`, return `'?'`.
134
135    precision: int, default 1
136        The number of decimal places to display for non-byte units.
137
138    Returns
139    -------
140    A human-readable string such as `'1.2 MB'` or `'340.0 kB'`.
141
142    Examples
143    --------
144    >>> format_bytes(0)
145    '0 B'
146    >>> format_bytes(1536)
147    '1.5 kB'
148    >>> format_bytes(None)
149    '?'
150    """
151    if num_bytes is None:
152        return '?'
153    try:
154        value = float(num_bytes)
155    except (TypeError, ValueError):
156        return '?'
157
158    sign = '-' if value < 0 else ''
159    value = abs(value)
160    units = ('B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB')
161    unit_index = 0
162    while value >= 1000.0 and unit_index < len(units) - 1:
163        value /= 1000.0
164        unit_index += 1
165
166    if unit_index == 0:
167        return f"{sign}{int(value)} {units[unit_index]}"
168    return f"{sign}{value:.{precision}f} {units[unit_index]}"
169
170def translate_rich_to_termcolor(*colors) -> tuple:
171    """Translate between rich and more_termcolor terminology."""
172    _colors = []
173    for c in colors:
174        _c_list = []
175        ### handle 'bright'
176        c = c.replace('bright_', 'bright ')
177
178        ### handle 'on'
179        if ' on ' in c:
180            _on = c.split(' on ')
181            _colors.append(_on[0])
182            for _c in _on[1:]:
183                _c_list.append('on ' + _c)
184        else:
185            _c_list += [c]
186
187        _colors += _c_list
188
189    return tuple(_colors)
190
191
192def rich_text_to_str(text: 'rich.text.Text', ansi: Optional[bool] = None) -> str:
193    """
194    Convert a `rich.text.Text` object to a string with ANSI in-tact.
195
196    Parameters
197    ----------
198    text: rich.text.Text
199        The text to render.
200
201    ansi: Optional[bool], default None
202        `True` renders the styles whatever `ANSI` says (text for another stream, such as a
203        warning for a terminal stderr while stdout is a file). Defaults to `get_console()`'s.
204    """
205    _console = _styled_console() if ansi else get_console()
206    if _console is None:
207        return str(text)
208    with _console.capture() as cap:
209        _console.print(text)
210    string = cap.get()
211    return string[:-1]
212
213
214_styled = None
215
216
217def _styled_console():
218    """A rich console that always renders styles (see `rich_text_to_str(ansi=True)`)."""
219    global _styled
220    if _styled is not None:
221        return _styled
222    from meerschaum.utils.packages import import_rich, attempt_import
223    import_rich(lazy=False)
224    rich_console = attempt_import('rich.console', lazy=False)
225    if rich_console is None:
226        return None
227    _styled = rich_console.Console(force_terminal=True, color_system='truecolor')
228    return _styled
229
230
231def _init():
232    """
233    Initial color settings (mostly for Windows).
234    """
235    if platform.system() != "Windows":
236        return
237    if 'PYTHONIOENCODING' not in os.environ:
238        os.environ['PYTHONIOENCODING'] = 'utf-8'
239    if 'PYTHONLEGACYWINDOWSSTDIO' not in os.environ:
240        os.environ['PYTHONLEGACYWINDOWSSTDIO'] = 'utf-8'
241    sys.stdin.reconfigure(encoding='utf-8')
242    sys.stdout.reconfigure(encoding='utf-8')
243    sys.stderr.reconfigure(encoding='utf-8')
244
245    from ctypes import windll
246    k = windll.kernel32
247    k.SetConsoleMode(k.GetStdHandle(-11), 7)
248    os.system("color")
249
250    from meerschaum.utils.packages import attempt_import
251    ### init colorama for Windows color output
252    colorama, more_termcolor = attempt_import(
253        'colorama',
254        'more_termcolor',
255        lazy = False,
256        warn = False,
257        color = False,
258    )
259    try:
260        colorama.init(autoreset=False)
261        success = True
262    except Exception:
263        import traceback
264        traceback.print_exc()
265        _attrs['ANSI'], _attrs['UNICODE'], _attrs['CHARSET'] = False, False, 'ascii'
266        success = False
267
268    if more_termcolor is None:
269        _attrs['ANSI'], _attrs['UNICODE'], _attrs['CHARSET'] = False, False, 'ascii'
270        success = False
271
272    return success
273
274_colorama_init = False
275def colored(text: str, *colors, as_rich_text: bool=False, **kw) -> Union[str, 'rich.text.Text']:
276    """Apply colors and rich styles to a string.
277    If a `style` keyword is provided, a `rich.text.Text` object will be parsed into a string.
278    Otherwise attempt to use the legacy `more_termcolor.colored` method.
279
280    Parameters
281    ----------
282    text: str
283        The string to apply formatting to.
284        
285    *colors:
286        A list of colors to pass to `more_termcolor.colored()`.
287        Has no effect if `style` is provided.
288
289    style: str, default None
290        If provided, pass to `rich` for processing.
291
292    as_rich_text: bool, default False
293        If `True`, return a `rich.Text` object.
294        `style` must be provided.
295        
296    **kw:
297        Keyword arguments to pass to `rich.text.Text` or `more_termcolor`.
298        
299
300    Returns
301    -------
302    An ANSI-formatted string or a `rich.text.Text` object if `as_rich_text` is `True`.
303
304    """
305    from meerschaum.utils.packages import import_rich, attempt_import
306    global _colorama_init
307    _init()
308    with _locks['_colorama_init']:
309        if not _colorama_init:
310            _colorama_init = _init()
311
312    if 'style' in kw:
313        import_rich(lazy=False)
314        rich_text = attempt_import('rich.text', lazy=False)
315        if rich_text is None:
316            ### `rich` is not a dependency of the base wheel: without it, text stays unstyled.
317            return text
318        text_obj = rich_text.Text(text, **kw)
319        if as_rich_text:
320            return text_obj
321        return rich_text_to_str(text_obj)
322
323    ### Off a terminal, under `NO_COLOR` or with `formatting:ansi` off, the text stays plain.
324    if not __getattr__('ANSI'):
325        return text
326
327    more_termcolor = attempt_import('more_termcolor', lazy=False)
328    try:
329        colored_text = more_termcolor.colored(text, *colors, **kw)
330    except Exception as e:
331        colored_text = None
332
333    if colored_text is not None:
334        return colored_text
335
336    try:
337        _colors = translate_rich_to_termcolor(*colors)
338        colored_text = more_termcolor.colored(text, *_colors, **kw)
339    except Exception as e:
340        colored_text = None
341
342    if colored_text is None:
343        ### NOTE: warn here?
344        return text
345
346    return colored_text
347
348console = None
349_console_ansi = None
350def get_console():
351    """
352    Get the rich console: styled (truecolor) when `ANSI` is on, plain text otherwise
353    (see `ansi_allowed()`).
354    """
355    global console, _console_ansi
356    ansi = bool(__getattr__('ANSI'))
357    if console is not None and _console_ansi == ansi:
358        return console
359    from meerschaum.utils.packages import import_rich, attempt_import
360    import_rich(lazy=False)
361    rich_console = attempt_import('rich.console', lazy=False)
362    if rich_console is None:
363        ### `rich` is not a dependency of the base wheel.
364        return None
365    try:
366        console = (
367            rich_console.Console(force_terminal=True, color_system='truecolor')
368            if ansi
369            else rich_console.Console(force_terminal=False, color_system=None, no_color=True)
370        )
371        _console_ansi = ansi
372    except Exception:
373        import traceback
374        traceback.print_exc()
375        console = None
376    return console
377
378
379def print_exception(**kw: Any) -> None:
380    """
381    Print the exception being handled: with `rich` when it is installed, else as plain text.
382    """
383    console = get_console()
384    if console is None:
385        import traceback
386        traceback.print_exc()
387        return
388    console.print_exception(**kw)
389
390
391def print_tuple(
392    tup: mrsm.SuccessTuple,
393    skip_common: bool = True,
394    common_only: bool = False,
395    upper_padding: int = 1,
396    lower_padding: int = 1,
397    left_padding: int = 1,
398    calm: bool = False,
399    _progress: Optional['rich.progress.Progress'] = None,
400) -> None:
401    """
402    Format `meerschaum.utils.typing.SuccessTuple`.
403
404    Parameters
405    ----------
406    skip_common: bool, default True
407        If `True`, do not print common success tuples (i.e. `(True, "Success")`).
408
409    common_only: bool, default False
410        If `True`, only print if the success tuple is common.
411
412    upper_padding: int, default 0
413        How many newlines to prepend to the message.
414
415    lower_padding: int, default 0
416        How many newlines to append to the message.
417
418    left_padding: int, default 1
419        How mant spaces to preprend to the message.
420
421    calm: bool, default False
422        If `True`, use the default emoji and color scheme.
423
424    """
425    from meerschaum._internal.static import STATIC_CONFIG
426    do_print = True
427
428    omit_messages = STATIC_CONFIG['system']['success']['ignore']
429
430    if common_only:
431        skip_common = False
432        do_print = tup[1].strip() in omit_messages
433
434    if skip_common:
435        do_print = tup[1].strip() not in omit_messages
436
437    if not do_print:
438        return
439
440    print(format_success_tuple(
441        tup,
442        upper_padding=upper_padding,
443        lower_padding=lower_padding,
444        calm=calm,
445        _progress=_progress,
446    ))
447
448
449def format_success_tuple(
450    tup: mrsm.SuccessTuple,
451    upper_padding: int = 0,
452    lower_padding: int = 0,
453    left_padding: int = 1,
454    calm: bool = False,
455    _progress: Optional['rich.progress.Progress'] = None,
456) -> str:
457    """
458    Format `meerschaum.utils.typing.SuccessTuple`.
459
460    Parameters
461    ----------
462    upper_padding: int, default 0
463        How many newlines to prepend to the message.
464
465    lower_padding: int, default 0
466        How many newlines to append to the message.
467
468    left_padding: int, default 1
469        How mant spaces to preprend to the message.
470
471    calm: bool, default False
472        If `True`, use the default emoji and color scheme.
473    """
474    _init()
475    try:
476        status = 'success' if tup[0] else 'failure'
477    except TypeError:
478        status = 'failure'
479        tup = None, None
480
481    if calm:
482        status += '_calm'
483
484    ANSI, CHARSET = __getattr__('ANSI'), __getattr__('CHARSET')
485    from meerschaum.config import get_config
486    status_config = get_config('formatting', status, patch=True)
487
488    msg = (' ' * left_padding) + status_config[CHARSET]['icon'] + ' ' + str(highlight_pipes(tup[1]))
489    lines = msg.split('\n')
490    lines = [lines[0]] + [
491        (('    ' + line if not line.startswith(' ') else line))
492        for line in lines[1:]
493    ]
494    if ANSI:
495        lines[0] = fill_ansi(lines[0], **status_config['ansi']['rich'])
496
497    msg = '\n'.join(lines)
498    msg = ('\n' * upper_padding) + msg + ('\n' * lower_padding)
499    return msg
500
501
502def print_options(
503    options: Optional[Iterable[Any]] = None,
504    nopretty: bool = False,
505    no_rich: bool = False,
506    name: str = 'options',
507    header: Optional[str] = None,
508    num_cols: Optional[int] = None,
509    adjust_cols: bool = True,
510    sort_options: bool = False,
511    number_options: bool = False,
512    **kw
513) -> None:
514    """
515    Print items in an iterable as a fancy table.
516
517    Parameters
518    ----------
519    options: Optional[Dict[str, Any]], default None
520        The iterable to be printed.
521
522    nopretty: bool, default False
523        If `True`, don't use fancy formatting.
524
525    no_rich: bool, default False
526        If `True`, don't use `rich` to format the output.
527
528    name: str, default 'options'
529        The text in the default header after `'Available'`.
530
531    header: Optional[str], default None
532        If provided, override `name` and use this as the header text.
533
534    num_cols: Optional[int], default None
535        How many columns in the table. Depends on the terminal size. If `None`, use 8.
536
537    adjust_cols: bool, default True
538        If `True`, adjust the number of columns depending on the terminal size.
539
540    sort_options: bool, default False
541        If `True`, print the options in sorted order.
542
543    number_options: bool, default False
544        If `True`, print the option's number in the list (1 index).
545
546    """
547    from meerschaum.utils.packages import import_rich
548    from meerschaum.utils.formatting import highlight_pipes
549    from meerschaum.utils.misc import get_cols_lines, string_width
550
551    if options is None:
552        options = {}
553    _options = []
554    for o in options:
555        _options.append(str(o))
556    if sort_options:
557        _options = sorted(_options)
558    _header = f"\nAvailable {name}" if header is None else header
559
560    if num_cols is None:
561        num_cols = 8
562
563    def _print_options_no_rich():
564        if not nopretty:
565            print()
566            print(make_header(_header))
567        ### print actions
568        for i, option in enumerate(_options):
569            marker = '-' if not number_options else (str(i + 1) + '.')
570            if not nopretty:
571                print(f"  {marker} ", end="")
572            print(option)
573        if not nopretty:
574            print()
575
576    ### `rich` is not a dependency of the base wheel: without it, options print as plain text.
577    rich = import_rich(lazy=False)
578    if rich is None or nopretty or no_rich or get_console() is None:
579        _print_options_no_rich()
580        return None
581
582    ### Prevent too many options from being truncated on small terminals.
583    if adjust_cols and _options:
584        _cols, _lines = get_cols_lines()
585        while num_cols > 1:
586            cell_len = int(((_cols - 4) - (3 * (num_cols - 1))) / num_cols)
587            num_too_big = sum([(1 if string_width(o) > cell_len else 0) for o in _options])
588            if num_too_big > int(len(_options) / 3):
589                num_cols -= 1
590                continue
591            break
592
593    from meerschaum.utils.packages import attempt_import
594    rich_table = attempt_import('rich.table')
595    Text = attempt_import('rich.text').Text
596    box = attempt_import('rich.box')
597    Table = rich_table.Table
598
599    if _header is not None:
600        table = Table(
601            title=_header,
602            box=box.SIMPLE,
603            show_header=False,
604            show_footer=False,
605            title_style='',
606            expand = True,
607        )
608    else:
609        table = Table.grid(padding=(0, 2))
610    for i in range(num_cols):
611        table.add_column()
612
613    if len(_options) < 12:
614        ### If fewer than 12 items, use a single column
615        for i, option in enumerate(_options):
616            item = highlight_pipes(option)
617            if number_options:
618                item = str(i + 1) + '. ' + item
619            table.add_row(Text.from_ansi(item))
620    else:
621        ### Otherwise, use multiple columns as before
622        num_rows = (len(_options) + num_cols - 1) // num_cols
623        item_ix = 0
624        for i in range(num_rows):
625            row = []
626            for j in range(num_cols):
627                index = i + j * num_rows
628                if index < len(_options):
629                    item = highlight_pipes(_options[index])
630                    if number_options:
631                        item = str(i + 1) + '. ' + item
632                    row.append(Text.from_ansi(item))
633                    item_ix += 1
634                else:
635                    row.append('')
636            table.add_row(*row)
637
638    get_console().print(table)
639    return None
640
641
642def fill_ansi(string: str, style: str = '', ansi: Optional[bool] = None) -> str:
643    """
644    Fill in non-formatted segments of ANSI text.
645
646    Parameters
647    ----------
648    string: str
649        A string which contains ANSI escape codes.
650
651    style: str
652        Style arguments to pass to `rich.text.Text`.
653
654    ansi: Optional[bool], default None
655        `True` keeps the styles whatever `ANSI` says (see `rich_text_to_str()`).
656
657    Returns
658    -------
659    A string with ANSI styling applied to the segments which don't yet have a style applied.
660    """
661    from meerschaum.utils.packages import import_rich, attempt_import
662    from meerschaum.utils.misc import iterate_chunks
663    _ = import_rich(lazy=False)
664    rich_ansi, rich_text = attempt_import('rich.ansi', 'rich.text', lazy=False)
665    if rich_ansi is None or rich_text is None:
666        ### `rich` is not a dependency of the base wheel: without it, text stays unstyled.
667        return string
668    Text = rich_text.Text
669    try:
670        msg = Text.from_ansi(string)
671    except AttributeError:
672        import traceback
673        traceback.print_exc()
674        msg = ''
675
676    plain_indices = []
677    for left_span, right_span in iterate_chunks(msg.spans, 2, fillvalue=len(msg)):
678        left = left_span.end
679        right = right_span.start if not isinstance(right_span, int) else right_span
680        if left != right:
681            plain_indices.append((left, right))
682    if msg.spans:
683        if msg.spans[0].start != 0:
684            plain_indices = [(0, msg.spans[0].start)] + plain_indices
685        if plain_indices and msg.spans[-1].end != len(msg) and plain_indices[-1][1] != len(msg):
686            plain_indices.append((msg.spans[-1].end, len(msg)))
687
688    if plain_indices:
689        for left, right in plain_indices:
690            msg.stylize(style, left, right)
691    else:
692        msg = Text(str(msg), style)
693
694    return rich_text_to_str(msg, ansi=ansi)
695
696
697def __getattr__(name: str) -> str:
698    """
699    Lazily load module-level variables.
700    """
701    if name.startswith('__') and name.endswith('__'):
702        raise AttributeError("Cannot import dunders from this module.")
703
704    if name in _attrs:
705        if _attrs[name] is None:
706            from meerschaum.config import get_config
707            if name.lower() in get_config('formatting'):
708                _attrs[name] = get_config('formatting', name.lower())
709            elif name == 'CHARSET':
710                _attrs[name] = 'unicode' if __getattr__('UNICODE') else 'ascii'
711        if name == 'ANSI':
712            ### `formatting:ansi`, and only where styles belong (`ansi_allowed()`).
713            return bool(_attrs[name]) and ansi_allowed()
714        return _attrs[name]
715    
716    if name == '__wrapped__':
717        import sys
718        return sys.modules[__name__]
719    if name == '__all__':
720        return __all__
721
722    try:
723        return globals()[name]
724    except KeyError:
725        raise AttributeError(f"Could not find '{name}'")