Pyspark: Предотвращение удаления терминальных символов ansi-escape при регистрации при использовании pyspark

Question

Я собрал пользовательский форматтер для регистратора, и я использую pyspark, но похоже, что весь мой цвет удаляется в командной строке. Я могу подтвердить, что escape-последовательности присутствуют в записи каждого испускаемого значения, но кажется, что они удаляются при отправке на терминал.

Почему?

import datetime 
import logging 

import colorama 
from pygments import highlight 
from pygments.lexers import JsonLexer 
from pygments.formatters import Terminal256Formatter 

# Required for colored output 
colorama.init() 


class CustomFormatter(logging.Formatter): 
    '''Modifies the level prefix of the log with the following level 
    information: 

    !!! - critical 
    ! - error 
    ? - warn 
     - info 
    - - debug 
    ''' 
    default_prefix = '???' # used with non-generic levels 

    color_mapping = { 
     logging.CRITICAL: colorama.Fore.RED + colorama.Style.BRIGHT, 
     logging.ERROR: colorama.Fore.RED + colorama.Style.BRIGHT, 
     logging.WARNING: colorama.Fore.YELLOW + colorama.Style.BRIGHT, 
     logging.DEBUG: colorama.Style.DIM, 
    } 

    prefix_mapping = { 
     logging.CRITICAL: '!!!', 
     logging.ERROR: ' ! ', 
     logging.WARNING: ' ? ', 
     logging.INFO: ' ', 
     logging.DEBUG: ' · ', 
    } 

    def format(self, record): 
     # Capture relevant record data 
     level = self.prefix_mapping.get(record.levelno) or self.default_prefix 
     msecs = datetime.datetime.fromtimestamp(record.created).strftime("%Y-%m-%d %H:%M:%S") 
     msg = record.msg.rstrip('\n') 

     # Setup colors 
     color = self.color_mapping.get(record.levelno) or '' 
     dim = colorama.Style.DIM 
     reset = colorama.Fore.RESET + colorama.Style.RESET_ALL 
     name = record.name 
     func = record.funcName 

     # Setup output 
     lexer = JsonLexer() 
     formatter = Terminal256Formatter() 
     try: 
      msg = '\n'.join(
       highlight(m, lexer, formatter).rstrip('\n') 
       for m in msg.split('\n') 
      ) 
     except: 
      pass 
     data = {k: v for k, v in locals().items()} 
     d = '{color}{level}{reset} {dim}{msecs} [{name}]{reset} {msg}'.format(**data) 
     record.msg = d 

     # Dump 
     return super(CustomFormatter, self).format(record) 

Использование:

import logging 
from CustomFormatter import CustomFormatter 

def get_logger(name, level=None): 
    level = logging.DEBUG if not isinstance(level, int) else level 
    handler = logging.StreamHandler(sys.stdout) 
    handler.level = level or logging.INFO 
    formatter = CustomFormatter() 
    handler.setFormatter(formatter) 
    logger = logging.getLogger(name) 
    logger.addHandler(handler) 
    return logger 

logger = get_logger('tester') 
logger.error('Error here') 

Answer 1

Я провел некоторое время углубиться в это, и я обнаружил, что последовательности терминальных побег различны при загрузке под pyspark. То, как я исправил это, - это использовать pygments для запуска над созданным мной вывода терминала (см. Функцию: fix_for_spark).

# -*- coding: utf-8 -*- 
import datetime 
import json 
import logging 
import os 

import colorama 
from pygments import highlight 
from pygments.lexers import JsonLexer 
from pygments.formatters import Terminal256Formatter 

# Required for colored output 
colorama.init() 


class CustomFormatter(logging.Formatter): 

    '''Modifies the level prefix of the log with the following level 
    information: 

    !!! - critical 
    ! - error 
    ? - warn 
     - info 
    - - debug 
    ''' 
    default_prefix = '???' # used with non-generic levels 

    PYGMENTS_LEXER = JsonLexer() 
    PYGMENTS_FORMATTER = Terminal256Formatter() 

    color_mapping = { 
     logging.CRITICAL: colorama.Fore.RED + colorama.Style.BRIGHT, 
     logging.ERROR: colorama.Fore.RED + colorama.Style.BRIGHT, 
     logging.WARNING: colorama.Fore.YELLOW + colorama.Style.BRIGHT, 
     logging.DEBUG: colorama.Style.DIM, 
    } 

    prefix_mapping = { 

     logging.CRITICAL: '!!!', 
     logging.ERROR: ' ! ', 
     logging.WARNING: ' ? ', 
     logging.INFO: ' ️ ', 
     logging.DEBUG: ' · ', 
    } 

    def fix_for_spark(self, string): 
     if os.environ.get('SPARK_ENV_LOADED'): 
      # Setup output 
      new_string = [] 
      for s in string.split('\n'): 
       s = highlight(string, self.PYGMENTS_LEXER, self.PYGMENTS_FORMATTER) 
       new_string.append(s.rstrip('\n')) 
      string = '\n'.join(new_string) 
     return string 

    def format(self, record): 
     # Capture relevant record data 
     data = dict(
      level=self.prefix_mapping.get(record.levelno) or self.default_prefix, 
      msecs=datetime.datetime.fromtimestamp(record.created).strftime("%Y-%m-%d %H:%M:%S"), 

      # Setup colors 
      color=self.color_mapping.get(record.levelno) or '', 
      dim=colorama.Style.DIM, 
      reset=colorama.Fore.RESET + colorama.Style.RESET_ALL, 
      name=record.name, 
      func=record.funcName, 
     ) 

     # Format msg 
     prefix = '{color}{level}{reset} {dim}{msecs}{reset} {color}[{name}]{reset}' 
     prefix = prefix.format(**data) 
     prefix = self.fix_for_spark(prefix) 
     msg = record.msg 
     if not isinstance(msg, str): 
      try: 
       msg = json.dumps(msg, indent=4, sort_keys=True) 
      except: 
       msg = str(msg) 
     dmsg = [] 
     for m in msg.split('\n'): 
      m = highlight(m, self.PYGMENTS_LEXER, self.PYGMENTS_FORMATTER).rstrip('\n') 
      m = self.fix_for_spark(m) 
      dmsg.append(m) 
     dmsg = '\n'.join(dmsg) 
     data.update(locals().items()) 
     template = prefix + ' {msg}' 
     record.msg = '\n'.join(template.format(msg=m) for m in dmsg.split('\n')) 
     # Dump 
     return super(CustomFormatter, self).format(record)