printer.py

import pandas as pd
from collections import namedtuple


class Printer(object):
    def __init__(self, field_names, print_format="table", persistent_file=None):
        assert print_format in ("table", "normal")

        self.field_names_ = field_names
        self.format_ = print_format
        self.records_ = []
        self.handlers_ = dict()
        self.str_lens_ = dict()
        self.title_printed_ = False

        if persistent_file is not None:
            self.csv_ = open(persistent_file, "a")
        else:
            self.csv_ = None

        self.Record = None

    def __def__(self):
        if self.csv_ is not None:
            self.csv_.close()

    def finish(self):
        err = f"{len(self.field_names_)} vs. {len(self.handlers_)}"
        assert len(self.field_names_) == len(self.handlers_), err
        err = f"{len(self.field_names_)} vs. {len(self.str_lens_)}"
        assert len(self.field_names_) == len(self.str_lens_), err
        for fname in self.field_names_:
            assert fname in self.handlers_, f"{fname} handler not register"
            assert fname in self.str_lens_, f"{fname} str_len not register"

        self.Record = namedtuple("Record", self.field_names_)
        # DEBUG(zwx):
        # dummy = self.Record(*(["-"] * len(self.field_names_)))
        # df = pd.DataFrame(dummy)
        # if self.persistent_file_ is not None:
        #     df.to_csv(self.persistent_file_, mode='a', header=True)

    def record(self, *args, **kwargs):
        assert self.Record is not None
        r = self.Record(*args, **kwargs)
        self.records_.append(r)

    def register_handler(self, field, handler):
        assert callable(handler)
        self.handlers_[field] = handler

    def register_str_len(self, field, str_len):
        assert isinstance(str_len, int)
        self.str_lens_[field] = str_len

    def reset_records(self):
        self.records_ = []

    def print_table_title(self):
        fields = ""
        sep = ""

        for fname in self.field_names_:
            str_len = self.str_lens_[fname]
            fields += "| {} ".format(fname.ljust(str_len))
            sep += f"| {'-' * str_len} "

        fields += "|"
        sep += "|"
        print(fields)
        print(sep)
        self.title_printed_ = True

    def reset_title_printed(self):
        self.title_printed_ = False

    def print(self):
        df = pd.DataFrame(self.records_)
        fields = []
        for fname in self.field_names_:
            assert fname in self.handlers_
            handler = self.handlers_[fname]
            field_value = handler(df[fname])
            fields.append(field_value)

        if self.format_ == "table":
            if not self.title_printed_:
                self.print_table_title()

            record = ""
            for i, str_len in enumerate(self.str_lens_.values()):
                record += "| {} ".format(str(fields[i]).ljust(str_len))

            record += "|"
            print(record)

        elif self.format_ == "normal":
            record = ""

            for i, fval in enumerate(fields):
                fname = self.field_names_[i]
                record += f"{fname}: {fval}, "

            print(record)
        else:
            raise ValueError

        if self.csv_ is not None:
            df.to_csv(self.csv_, header=False)

        self.reset_records()