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    if console is not None:
256        return console
257    from meerschaum.utils.packages import import_rich, attempt_import
258    rich = import_rich()
259    rich_console = attempt_import('rich.console')
260    try:
261        console = rich_console.Console(force_terminal=True, color_system='truecolor')
262    except Exception:
263        import traceback
264        traceback.print_exc()
265        console = None
266    return console
267
268
269def print_tuple(
270    tup: mrsm.SuccessTuple,
271    skip_common: bool = True,
272    common_only: bool = False,
273    upper_padding: int = 1,
274    lower_padding: int = 1,
275    left_padding: int = 1,
276    calm: bool = False,
277    _progress: Optional['rich.progress.Progress'] = None,
278) -> None:
279    """
280    Format `meerschaum.utils.typing.SuccessTuple`.
281
282    Parameters
283    ----------
284    skip_common: bool, default True
285        If `True`, do not print common success tuples (i.e. `(True, "Success")`).
286
287    common_only: bool, default False
288        If `True`, only print if the success tuple is common.
289
290    upper_padding: int, default 0
291        How many newlines to prepend to the message.
292
293    lower_padding: int, default 0
294        How many newlines to append to the message.
295
296    left_padding: int, default 1
297        How mant spaces to preprend to the message.
298
299    calm: bool, default False
300        If `True`, use the default emoji and color scheme.
301
302    """
303    from meerschaum._internal.static import STATIC_CONFIG
304    do_print = True
305
306    omit_messages = STATIC_CONFIG['system']['success']['ignore']
307
308    if common_only:
309        skip_common = False
310        do_print = tup[1].strip() in omit_messages
311
312    if skip_common:
313        do_print = tup[1].strip() not in omit_messages
314
315    if not do_print:
316        return
317
318    print(format_success_tuple(
319        tup,
320        upper_padding=upper_padding,
321        lower_padding=lower_padding,
322        calm=calm,
323        _progress=_progress,
324    ))
325
326
327def format_success_tuple(
328    tup: mrsm.SuccessTuple,
329    upper_padding: int = 0,
330    lower_padding: int = 0,
331    left_padding: int = 1,
332    calm: bool = False,
333    _progress: Optional['rich.progress.Progress'] = None,
334) -> str:
335    """
336    Format `meerschaum.utils.typing.SuccessTuple`.
337
338    Parameters
339    ----------
340    upper_padding: int, default 0
341        How many newlines to prepend to the message.
342
343    lower_padding: int, default 0
344        How many newlines to append to the message.
345
346    left_padding: int, default 1
347        How mant spaces to preprend to the message.
348
349    calm: bool, default False
350        If `True`, use the default emoji and color scheme.
351    """
352    _init()
353    try:
354        status = 'success' if tup[0] else 'failure'
355    except TypeError:
356        status = 'failure'
357        tup = None, None
358
359    if calm:
360        status += '_calm'
361
362    ANSI, CHARSET = __getattr__('ANSI'), __getattr__('CHARSET')
363    from meerschaum.config import get_config
364    status_config = get_config('formatting', status, patch=True)
365
366    msg = (' ' * left_padding) + status_config[CHARSET]['icon'] + ' ' + str(highlight_pipes(tup[1]))
367    lines = msg.split('\n')
368    lines = [lines[0]] + [
369        (('    ' + line if not line.startswith(' ') else line))
370        for line in lines[1:]
371    ]
372    if ANSI:
373        lines[0] = fill_ansi(lines[0], **status_config['ansi']['rich'])
374
375    msg = '\n'.join(lines)
376    msg = ('\n' * upper_padding) + msg + ('\n' * lower_padding)
377    return msg
378
379
380def print_options(
381    options: Optional[Iterable[Any]] = None,
382    nopretty: bool = False,
383    no_rich: bool = False,
384    name: str = 'options',
385    header: Optional[str] = None,
386    num_cols: Optional[int] = None,
387    adjust_cols: bool = True,
388    sort_options: bool = False,
389    number_options: bool = False,
390    **kw
391) -> None:
392    """
393    Print items in an iterable as a fancy table.
394
395    Parameters
396    ----------
397    options: Optional[Dict[str, Any]], default None
398        The iterable to be printed.
399
400    nopretty: bool, default False
401        If `True`, don't use fancy formatting.
402
403    no_rich: bool, default False
404        If `True`, don't use `rich` to format the output.
405
406    name: str, default 'options'
407        The text in the default header after `'Available'`.
408
409    header: Optional[str], default None
410        If provided, override `name` and use this as the header text.
411
412    num_cols: Optional[int], default None
413        How many columns in the table. Depends on the terminal size. If `None`, use 8.
414
415    adjust_cols: bool, default True
416        If `True`, adjust the number of columns depending on the terminal size.
417
418    sort_options: bool, default False
419        If `True`, print the options in sorted order.
420
421    number_options: bool, default False
422        If `True`, print the option's number in the list (1 index).
423
424    """
425    from meerschaum.utils.packages import import_rich
426    from meerschaum.utils.formatting import highlight_pipes
427    from meerschaum.utils.misc import get_cols_lines, string_width
428
429    if options is None:
430        options = {}
431    _options = []
432    for o in options:
433        _options.append(str(o))
434    if sort_options:
435        _options = sorted(_options)
436    _header = f"\nAvailable {name}" if header is None else header
437
438    if num_cols is None:
439        num_cols = 8
440
441    def _print_options_no_rich():
442        if not nopretty:
443            print()
444            print(make_header(_header))
445        ### print actions
446        for i, option in enumerate(_options):
447            marker = '-' if not number_options else (str(i + 1) + '.')
448            if not nopretty:
449                print(f"  {marker} ", end="")
450            print(option)
451        if not nopretty:
452            print()
453
454    rich = import_rich()
455    if rich is None or nopretty or no_rich:
456        _print_options_no_rich()
457        return None
458
459    ### Prevent too many options from being truncated on small terminals.
460    if adjust_cols and _options:
461        _cols, _lines = get_cols_lines()
462        while num_cols > 1:
463            cell_len = int(((_cols - 4) - (3 * (num_cols - 1))) / num_cols)
464            num_too_big = sum([(1 if string_width(o) > cell_len else 0) for o in _options])
465            if num_too_big > int(len(_options) / 3):
466                num_cols -= 1
467                continue
468            break
469
470    from meerschaum.utils.packages import attempt_import
471    rich_table = attempt_import('rich.table')
472    Text = attempt_import('rich.text').Text
473    box = attempt_import('rich.box')
474    Table = rich_table.Table
475
476    if _header is not None:
477        table = Table(
478            title=_header,
479            box=box.SIMPLE,
480            show_header=False,
481            show_footer=False,
482            title_style='',
483            expand = True,
484        )
485    else:
486        table = Table.grid(padding=(0, 2))
487    for i in range(num_cols):
488        table.add_column()
489
490    if len(_options) < 12:
491        ### If fewer than 12 items, use a single column
492        for i, option in enumerate(_options):
493            item = highlight_pipes(option)
494            if number_options:
495                item = str(i + 1) + '. ' + item
496            table.add_row(Text.from_ansi(item))
497    else:
498        ### Otherwise, use multiple columns as before
499        num_rows = (len(_options) + num_cols - 1) // num_cols
500        item_ix = 0
501        for i in range(num_rows):
502            row = []
503            for j in range(num_cols):
504                index = i + j * num_rows
505                if index < len(_options):
506                    item = highlight_pipes(_options[index])
507                    if number_options:
508                        item = str(i + 1) + '. ' + item
509                    row.append(Text.from_ansi(item))
510                    item_ix += 1
511                else:
512                    row.append('')
513            table.add_row(*row)
514
515    get_console().print(table)
516    return None
517
518
519def fill_ansi(string: str, style: str = '') -> str:
520    """
521    Fill in non-formatted segments of ANSI text.
522
523    Parameters
524    ----------
525    string: str
526        A string which contains ANSI escape codes.
527
528    style: str
529        Style arguments to pass to `rich.text.Text`.
530
531    Returns
532    -------
533    A string with ANSI styling applied to the segments which don't yet have a style applied.
534    """
535    from meerschaum.utils.packages import import_rich, attempt_import
536    from meerschaum.utils.misc import iterate_chunks
537    _ = import_rich()
538    rich_ansi, rich_text = attempt_import('rich.ansi', 'rich.text')
539    Text = rich_text.Text
540    try:
541        msg = Text.from_ansi(string)
542    except AttributeError:
543        import traceback
544        traceback.print_exc()
545        msg = ''
546
547    plain_indices = []
548    for left_span, right_span in iterate_chunks(msg.spans, 2, fillvalue=len(msg)):
549        left = left_span.end
550        right = right_span.start if not isinstance(right_span, int) else right_span
551        if left != right:
552            plain_indices.append((left, right))
553    if msg.spans:
554        if msg.spans[0].start != 0:
555            plain_indices = [(0, msg.spans[0].start)] + plain_indices
556        if plain_indices and msg.spans[-1].end != len(msg) and plain_indices[-1][1] != len(msg):
557            plain_indices.append((msg.spans[-1].end, len(msg)))
558
559    if plain_indices:
560        for left, right in plain_indices:
561            msg.stylize(style, left, right)
562    else:
563        msg = Text(str(msg), style)
564
565    return rich_text_to_str(msg)
566
567
568def __getattr__(name: str) -> str:
569    """
570    Lazily load module-level variables.
571    """
572    if name.startswith('__') and name.endswith('__'):
573        raise AttributeError("Cannot import dunders from this module.")
574
575    if name in _attrs:
576        if _attrs[name] is not None:
577            return _attrs[name]
578        from meerschaum.config import get_config
579        if name.lower() in get_config('formatting'):
580            _attrs[name] = get_config('formatting', name.lower())
581        elif name == 'CHARSET':
582            _attrs[name] = 'unicode' if __getattr__('UNICODE') else 'ascii'
583        return _attrs[name]
584    
585    if name == '__wrapped__':
586        import sys
587        return sys.modules[__name__]
588    if name == '__all__':
589        return __all__
590
591    try:
592        return globals()[name]
593    except KeyError:
594        raise AttributeError(f"Could not find '{name}'")