File size: 2,106 Bytes
4943752 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
"""
Attack Logs to file
========================
"""
import os
import sys
import terminaltables
from textattack.shared import logger
from .logger import Logger
class FileLogger(Logger):
"""Logs the results of an attack to a file, or `stdout`."""
def __init__(self, filename="", stdout=False, color_method="ansi"):
self.stdout = stdout
self.filename = filename
self.color_method = color_method
if stdout:
self.fout = sys.stdout
elif isinstance(filename, str):
directory = os.path.dirname(filename)
directory = directory if directory else "."
if not os.path.exists(directory):
os.makedirs(directory)
self.fout = open(filename, "w")
logger.info(f"Logging to text file at path {filename}")
else:
self.fout = filename
self.num_results = 0
def __getstate__(self):
# Temporarily save file handle b/c we can't copy it
state = {i: self.__dict__[i] for i in self.__dict__ if i != "fout"}
return state
def __setstate__(self, state):
self.__dict__ = state
if self.stdout:
self.fout = sys.stdout
else:
self.fout = open(self.filename, "a")
def log_attack_result(self, result):
self.num_results += 1
# if self.stdout and sys.stdout.isatty():
self.fout.write(
"-" * 45 + " Result " + str(self.num_results) + " " + "-" * 45 + "\n"
)
self.fout.write(result.__str__(color_method=self.color_method))
self.fout.write("\n")
def log_summary_rows(self, rows, title, window_id):
if self.stdout:
table_rows = [[title, ""]] + rows
table = terminaltables.AsciiTable(table_rows)
self.fout.write(table.table)
else:
for row in rows:
self.fout.write(f"{row[0]} {row[1]}\n")
def log_sep(self):
self.fout.write("-" * 90 + "\n")
def flush(self):
self.fout.flush()
def close(self):
self.fout.close()
|