diff options
| -rw-r--r-- | cli.py | 6 | ||||
| -rw-r--r-- | obfuscator/write.py | 24 | ||||
| -rw-r--r-- | test/test_write.py | 2 |
3 files changed, 19 insertions, 13 deletions
@@ -1,8 +1,8 @@ import argparse from obfuscator.read import DataReader +from obfuscator.write import DataWriter from obfuscator.obfuscate import obfuscate from obfuscator.logger import get_logger -from obfuscator.csv_writer import create_byte_stream def main(): @@ -43,7 +43,9 @@ def main(): obfuscated_data = obfuscate(data, args.pii) - return create_byte_stream(obfuscated_data) + writer = DataWriter() + + return writer.create_byte_stream(obfuscated_data) if __name__ == "__main__": diff --git a/obfuscator/write.py b/obfuscator/write.py index de7cd4b..4081f0f 100644 --- a/obfuscator/write.py +++ b/obfuscator/write.py @@ -6,18 +6,22 @@ from obfuscator.logger import get_logger logger = get_logger("CSVWRITER") -def create_byte_stream(data: List[Dict[str, str]]) -> bytes: - if not data: - logger.error("Invalid or empty data was provided to write") +class DataWriter: + def __init__(self): + pass - output = io.StringIO() + def create_byte_stream(self, data: List[Dict[str, str]]) -> bytes: + if not data: + logger.error("Invalid or empty data was provided to write") - headers = list(data[0].keys()) + output = io.StringIO() - writer = csv.DictWriter(output, fieldnames=headers) - writer.writeheader() - writer.writerows(data) + headers = list(data[0].keys()) - csv_string = output.getvalue() + writer = csv.DictWriter(output, fieldnames=headers) + writer.writeheader() + writer.writerows(data) - return csv_string.encode("utf-8") + csv_string = output.getvalue() + + return csv_string.encode("utf-8") diff --git a/test/test_write.py b/test/test_write.py index eceac28..f339799 100644 --- a/test/test_write.py +++ b/test/test_write.py @@ -1,6 +1,6 @@ import io import csv -from obfuscator.csv_writer import create_byte_stream +from obfuscator.write import create_byte_stream def csv_bytes_to_list(csv_bytes: bytes): |