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    'colored',
 37    'translate_rich_to_termcolor',
 38    'get_console',
 39    'format_success_tuple',
 40    'print_tuple',
 41    'print_options',
 42    'fill_ansi',
 43    'pprint',
 44    'pprint_df',
 45    'format_dataframe',
 46    'highlight_pipes',
 47    'pprint_pipes',
 48    'make_header',
 49    'pipe_repr',
 50    'print_pipes_results',
 51    'extract_stats_from_message',
 52    'format_bytes',
 53])
 54__pdoc__ = {}
 55_locks = {
 56    '_colorama_init': RLock(),
 57}
 58
 59
 60def colored_fallback(*args, **kw):
 61    return ' '.join(args)
 62
 63
 64def format_bytes(num_bytes: Optional[Union[int, float]], precision: int = 1) -> str:
 65    """
 66    Return a human-readable representation of a number of bytes.
 67
 68    Parameters
 69    ----------
 70    num_bytes: Optional[Union[int, float]]
 71        The number of bytes to format. If `None`, return `'?'`.
 72
 73    precision: int, default 1
 74        The number of decimal places to display for non-byte units.
 75
 76    Returns
 77    -------
 78    A human-readable string such as `'1.2 MB'` or `'340.0 kB'`.
 79
 80    Examples
 81    --------
 82    >>> format_bytes(0)
 83    '0 B'
 84    >>> format_bytes(1536)
 85    '1.5 kB'
 86    >>> format_bytes(None)
 87    '?'
 88    """
 89    if num_bytes is None:
 90        return '?'
 91    try:
 92        value = float(num_bytes)
 93    except (TypeError, ValueError):
 94        return '?'
 95
 96    sign = '-' if value < 0 else ''
 97    value = abs(value)
 98    units = ('B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB')
 99    unit_index = 0
100    while value >= 1000.0 and unit_index < len(units) - 1:
101        value /= 1000.0
102        unit_index += 1
103
104    if unit_index == 0:
105        return f"{sign}{int(value)} {units[unit_index]}"
106    return f"{sign}{value:.{precision}f} {units[unit_index]}"
107
108def translate_rich_to_termcolor(*colors) -> tuple:
109    """Translate between rich and more_termcolor terminology."""
110    _colors = []
111    for c in colors:
112        _c_list = []
113        ### handle 'bright'
114        c = c.replace('bright_', 'bright ')
115
116        ### handle 'on'
117        if ' on ' in c:
118            _on = c.split(' on ')
119            _colors.append(_on[0])
120            for _c in _on[1:]:
121                _c_list.append('on ' + _c)
122        else:
123            _c_list += [c]
124
125        _colors += _c_list
126
127    return tuple(_colors)
128
129
130def rich_text_to_str(text: 'rich.text.Text') -> str:
131    """Convert a `rich.text.Text` object to a string with ANSI in-tact."""
132    _console = get_console()
133    if _console is None:
134        return str(text)
135    with console.capture() as cap:
136        console.print(text)
137    string = cap.get()
138    return string[:-1]
139
140
141def _init():
142    """
143    Initial color settings (mostly for Windows).
144    """
145    if platform.system() != "Windows":
146        return
147    if 'PYTHONIOENCODING' not in os.environ:
148        os.environ['PYTHONIOENCODING'] = 'utf-8'
149    if 'PYTHONLEGACYWINDOWSSTDIO' not in os.environ:
150        os.environ['PYTHONLEGACYWINDOWSSTDIO'] = 'utf-8'
151    sys.stdin.reconfigure(encoding='utf-8')
152    sys.stdout.reconfigure(encoding='utf-8')
153    sys.stderr.reconfigure(encoding='utf-8')
154
155    from ctypes import windll
156    k = windll.kernel32
157    k.SetConsoleMode(k.GetStdHandle(-11), 7)
158    os.system("color")
159
160    from meerschaum.utils.packages import attempt_import
161    ### init colorama for Windows color output
162    colorama, more_termcolor = attempt_import(
163        'colorama',
164        'more_termcolor',
165        lazy = False,
166        warn = False,
167        color = False,
168    )
169    try:
170        colorama.init(autoreset=False)
171        success = True
172    except Exception:
173        import traceback
174        traceback.print_exc()
175        _attrs['ANSI'], _attrs['UNICODE'], _attrs['CHARSET'] = False, False, 'ascii'
176        success = False
177
178    if more_termcolor is None:
179        _attrs['ANSI'], _attrs['UNICODE'], _attrs['CHARSET'] = False, False, 'ascii'
180        success = False
181
182    return success
183
184_colorama_init = False
185def colored(text: str, *colors, as_rich_text: bool=False, **kw) -> Union[str, 'rich.text.Text']:
186    """Apply colors and rich styles to a string.
187    If a `style` keyword is provided, a `rich.text.Text` object will be parsed into a string.
188    Otherwise attempt to use the legacy `more_termcolor.colored` method.
189
190    Parameters
191    ----------
192    text: str
193        The string to apply formatting to.
194        
195    *colors:
196        A list of colors to pass to `more_termcolor.colored()`.
197        Has no effect if `style` is provided.
198
199    style: str, default None
200        If provided, pass to `rich` for processing.
201
202    as_rich_text: bool, default False
203        If `True`, return a `rich.Text` object.
204        `style` must be provided.
205        
206    **kw:
207        Keyword arguments to pass to `rich.text.Text` or `more_termcolor`.
208        
209
210    Returns
211    -------
212    An ANSI-formatted string or a `rich.text.Text` object if `as_rich_text` is `True`.
213
214    """
215    from meerschaum.utils.packages import import_rich, attempt_import
216    global _colorama_init
217    _init()
218    with _locks['_colorama_init']:
219        if not _colorama_init:
220            _colorama_init = _init()
221
222    if 'style' in kw:
223        rich = import_rich()
224        rich_text = attempt_import('rich.text')
225        text_obj = rich_text.Text(text, **kw)
226        if as_rich_text:
227            return text_obj
228        return rich_text_to_str(text_obj)
229
230    more_termcolor = attempt_import('more_termcolor', lazy=False)
231    try:
232        colored_text = more_termcolor.colored(text, *colors, **kw)
233    except Exception as e:
234        colored_text = None
235
236    if colored_text is not None:
237        return colored_text
238
239    try:
240        _colors = translate_rich_to_termcolor(*colors)
241        colored_text = more_termcolor.colored(text, *_colors, **kw)
242    except Exception as e:
243        colored_text = None
244
245    if colored_text is None:
246        ### NOTE: warn here?
247        return text
248
249    return colored_text
250
251console = None
252def get_console():
253    """Get the rich console."""
254    global console
255    from meerschaum.utils.packages import import_rich, attempt_import
256    rich = import_rich()
257    rich_console = attempt_import('rich.console')
258    try:
259        console = rich_console.Console(force_terminal=True, color_system='truecolor')
260    except Exception:
261        import traceback
262        traceback.print_exc()
263        console = None
264    return console
265
266
267def print_tuple(
268    tup: mrsm.SuccessTuple,
269    skip_common: bool = True,
270    common_only: bool = False,
271    upper_padding: int = 1,
272    lower_padding: int = 1,
273    left_padding: int = 1,
274    calm: bool = False,
275    _progress: Optional['rich.progress.Progress'] = None,
276) -> None:
277    """
278    Format `meerschaum.utils.typing.SuccessTuple`.
279
280    Parameters
281    ----------
282    skip_common: bool, default True
283        If `True`, do not print common success tuples (i.e. `(True, "Success")`).
284
285    common_only: bool, default False
286        If `True`, only print if the success tuple is common.
287
288    upper_padding: int, default 0
289        How many newlines to prepend to the message.
290
291    lower_padding: int, default 0
292        How many newlines to append to the message.
293
294    left_padding: int, default 1
295        How mant spaces to preprend to the message.
296
297    calm: bool, default False
298        If `True`, use the default emoji and color scheme.
299
300    """
301    from meerschaum._internal.static import STATIC_CONFIG
302    do_print = True
303
304    omit_messages = STATIC_CONFIG['system']['success']['ignore']
305
306    if common_only:
307        skip_common = False
308        do_print = tup[1].strip() in omit_messages
309
310    if skip_common:
311        do_print = tup[1].strip() not in omit_messages
312
313    if not do_print:
314        return
315
316    print(format_success_tuple(
317        tup,
318        upper_padding=upper_padding,
319        lower_padding=lower_padding,
320        calm=calm,
321        _progress=_progress,
322    ))
323
324
325def format_success_tuple(
326    tup: mrsm.SuccessTuple,
327    upper_padding: int = 0,
328    lower_padding: int = 0,
329    left_padding: int = 1,
330    calm: bool = False,
331    _progress: Optional['rich.progress.Progress'] = None,
332) -> str:
333    """
334    Format `meerschaum.utils.typing.SuccessTuple`.
335
336    Parameters
337    ----------
338    upper_padding: int, default 0
339        How many newlines to prepend to the message.
340
341    lower_padding: int, default 0
342        How many newlines to append to the message.
343
344    left_padding: int, default 1
345        How mant spaces to preprend to the message.
346
347    calm: bool, default False
348        If `True`, use the default emoji and color scheme.
349    """
350    _init()
351    try:
352        status = 'success' if tup[0] else 'failure'
353    except TypeError:
354        status = 'failure'
355        tup = None, None
356
357    if calm:
358        status += '_calm'
359
360    ANSI, CHARSET = __getattr__('ANSI'), __getattr__('CHARSET')
361    from meerschaum.config import get_config
362    status_config = get_config('formatting', status, patch=True)
363
364    msg = (' ' * left_padding) + status_config[CHARSET]['icon'] + ' ' + str(highlight_pipes(tup[1]))
365    lines = msg.split('\n')
366    lines = [lines[0]] + [
367        (('    ' + line if not line.startswith(' ') else line))
368        for line in lines[1:]
369    ]
370    if ANSI:
371        lines[0] = fill_ansi(lines[0], **status_config['ansi']['rich'])
372
373    msg = '\n'.join(lines)
374    msg = ('\n' * upper_padding) + msg + ('\n' * lower_padding)
375    return msg
376
377
378def print_options(
379    options: Optional[Iterable[Any]] = None,
380    nopretty: bool = False,
381    no_rich: bool = False,
382    name: str = 'options',
383    header: Optional[str] = None,
384    num_cols: Optional[int] = None,
385    adjust_cols: bool = True,
386    sort_options: bool = False,
387    number_options: bool = False,
388    **kw
389) -> None:
390    """
391    Print items in an iterable as a fancy table.
392
393    Parameters
394    ----------
395    options: Optional[Dict[str, Any]], default None
396        The iterable to be printed.
397
398    nopretty: bool, default False
399        If `True`, don't use fancy formatting.
400
401    no_rich: bool, default False
402        If `True`, don't use `rich` to format the output.
403
404    name: str, default 'options'
405        The text in the default header after `'Available'`.
406
407    header: Optional[str], default None
408        If provided, override `name` and use this as the header text.
409
410    num_cols: Optional[int], default None
411        How many columns in the table. Depends on the terminal size. If `None`, use 8.
412
413    adjust_cols: bool, default True
414        If `True`, adjust the number of columns depending on the terminal size.
415
416    sort_options: bool, default False
417        If `True`, print the options in sorted order.
418
419    number_options: bool, default False
420        If `True`, print the option's number in the list (1 index).
421
422    """
423    from meerschaum.utils.packages import import_rich
424    from meerschaum.utils.formatting import highlight_pipes
425    from meerschaum.utils.misc import get_cols_lines, string_width
426
427    if options is None:
428        options = {}
429    _options = []
430    for o in options:
431        _options.append(str(o))
432    if sort_options:
433        _options = sorted(_options)
434    _header = f"\nAvailable {name}" if header is None else header
435
436    if num_cols is None:
437        num_cols = 8
438
439    def _print_options_no_rich():
440        if not nopretty:
441            print()
442            print(make_header(_header))
443        ### print actions
444        for i, option in enumerate(_options):
445            marker = '-' if not number_options else (str(i + 1) + '.')
446            if not nopretty:
447                print(f"  {marker} ", end="")
448            print(option)
449        if not nopretty:
450            print()
451
452    rich = import_rich()
453    if rich is None or nopretty or no_rich:
454        _print_options_no_rich()
455        return None
456
457    ### Prevent too many options from being truncated on small terminals.
458    if adjust_cols and _options:
459        _cols, _lines = get_cols_lines()
460        while num_cols > 1:
461            cell_len = int(((_cols - 4) - (3 * (num_cols - 1))) / num_cols)
462            num_too_big = sum([(1 if string_width(o) > cell_len else 0) for o in _options])
463            if num_too_big > int(len(_options) / 3):
464                num_cols -= 1
465                continue
466            break
467
468    from meerschaum.utils.packages import attempt_import
469    rich_table = attempt_import('rich.table')
470    Text = attempt_import('rich.text').Text
471    box = attempt_import('rich.box')
472    Table = rich_table.Table
473
474    if _header is not None:
475        table = Table(
476            title=_header,
477            box=box.SIMPLE,
478            show_header=False,
479            show_footer=False,
480            title_style='',
481            expand = True,
482        )
483    else:
484        table = Table.grid(padding=(0, 2))
485    for i in range(num_cols):
486        table.add_column()
487
488    if len(_options) < 12:
489        ### If fewer than 12 items, use a single column
490        for i, option in enumerate(_options):
491            item = highlight_pipes(option)
492            if number_options:
493                item = str(i + 1) + '. ' + item
494            table.add_row(Text.from_ansi(item))
495    else:
496        ### Otherwise, use multiple columns as before
497        num_rows = (len(_options) + num_cols - 1) // num_cols
498        item_ix = 0
499        for i in range(num_rows):
500            row = []
501            for j in range(num_cols):
502                index = i + j * num_rows
503                if index < len(_options):
504                    item = highlight_pipes(_options[index])
505                    if number_options:
506                        item = str(i + 1) + '. ' + item
507                    row.append(Text.from_ansi(item))
508                    item_ix += 1
509                else:
510                    row.append('')
511            table.add_row(*row)
512
513    get_console().print(table)
514    return None
515
516
517def fill_ansi(string: str, style: str = '') -> str:
518    """
519    Fill in non-formatted segments of ANSI text.
520
521    Parameters
522    ----------
523    string: str
524        A string which contains ANSI escape codes.
525
526    style: str
527        Style arguments to pass to `rich.text.Text`.
528
529    Returns
530    -------
531    A string with ANSI styling applied to the segments which don't yet have a style applied.
532    """
533    from meerschaum.utils.packages import import_rich, attempt_import
534    from meerschaum.utils.misc import iterate_chunks
535    _ = import_rich()
536    rich_ansi, rich_text = attempt_import('rich.ansi', 'rich.text')
537    Text = rich_text.Text
538    try:
539        msg = Text.from_ansi(string)
540    except AttributeError:
541        import traceback
542        traceback.print_exc()
543        msg = ''
544
545    plain_indices = []
546    for left_span, right_span in iterate_chunks(msg.spans, 2, fillvalue=len(msg)):
547        left = left_span.end
548        right = right_span.start if not isinstance(right_span, int) else right_span
549        if left != right:
550            plain_indices.append((left, right))
551    if msg.spans:
552        if msg.spans[0].start != 0:
553            plain_indices = [(0, msg.spans[0].start)] + plain_indices
554        if plain_indices and msg.spans[-1].end != len(msg) and plain_indices[-1][1] != len(msg):
555            plain_indices.append((msg.spans[-1].end, len(msg)))
556
557    if plain_indices:
558        for left, right in plain_indices:
559            msg.stylize(style, left, right)
560    else:
561        msg = Text(str(msg), style)
562
563    return rich_text_to_str(msg)
564
565
566def __getattr__(name: str) -> str:
567    """
568    Lazily load module-level variables.
569    """
570    if name.startswith('__') and name.endswith('__'):
571        raise AttributeError("Cannot import dunders from this module.")
572
573    if name in _attrs:
574        if _attrs[name] is not None:
575            return _attrs[name]
576        from meerschaum.config import get_config
577        if name.lower() in get_config('formatting'):
578            _attrs[name] = get_config('formatting', name.lower())
579        elif name == 'CHARSET':
580            _attrs[name] = 'unicode' if __getattr__('UNICODE') else 'ascii'
581        return _attrs[name]
582    
583    if name == '__wrapped__':
584        import sys
585        return sys.modules[__name__]
586    if name == '__all__':
587        return __all__
588
589    try:
590        return globals()[name]
591    except KeyError:
592        raise AttributeError(f"Could not find '{name}'")