codekingpro/portable-devtools

1"""2Welcome to the encoding dance!3 4In a nutshell, text columns are actually a proxy class for byte columns,5which just encode/decodes contents.6"""7 8from mitmproxy.tools.console import signals9from mitmproxy.tools.console.grideditor import col_bytes10 11 12class Column(col_bytes.Column):13    def __init__(self, heading, encoding="utf8", errors="surrogateescape"):14        super().__init__(heading)15        self.encoding_args = encoding, errors16 17    def Display(self, data):18        return TDisplay(data, self.encoding_args)19 20    def Edit(self, data):21        return TEdit(data, self.encoding_args)22 23    def blank(self):24        return ""25 26 27# This is the same for both edit and display.28class EncodingMixin:29    def __init__(self, data, encoding_args):30        self.encoding_args = encoding_args31        super().__init__(str(data).encode(*self.encoding_args))  # type: ignore32 33    def get_data(self):34        data = super().get_data()  # type: ignore35        try:36            return data.decode(*self.encoding_args)37        except ValueError:38            signals.status_message.send(message="Invalid encoding.")39            raise40 41 42# urwid forces a different name for a subclass.43class TDisplay(EncodingMixin, col_bytes.Display):44    pass45 46 47class TEdit(EncodingMixin, col_bytes.Edit):48    pass49