diff options
| -rw-r--r-- | vk/tracking/db/backend/csv.py | 5 | ||||
| -rw-r--r-- | vk/tracking/db/format.py | 4 | ||||
| -rw-r--r-- | vk/tracking/db/io.py | 56 | ||||
| -rw-r--r-- | vk/utils/io.py | 27 |
4 files changed, 30 insertions, 62 deletions
diff --git a/vk/tracking/db/backend/csv.py b/vk/tracking/db/backend/csv.py index 43038e4..0503872 100644 --- a/vk/tracking/db/backend/csv.py +++ b/vk/tracking/db/backend/csv.py @@ -3,15 +3,16 @@ # For details, see https://github.com/egor-tensin/vk-scripts. # Distributed under the MIT License. +from vk.utils.io import FileReaderCSV, FileWriterCSV + from .. import meta -from ..io import FileReaderCSV, FileWriterCSV from ..record import Record from ..timestamp import Timestamp class Writer(meta.Writer): def __init__(self, fd): - self._writer = FileWriterCSV(fd) + self._writer = FileWriterCSV(fd, flush=True) def on_initial_status(self, user): self._write_record(user) diff --git a/vk/tracking/db/format.py b/vk/tracking/db/format.py index 028d403..2a54f6d 100644 --- a/vk/tracking/db/format.py +++ b/vk/tracking/db/format.py @@ -6,7 +6,9 @@ from enum import Enum import sys -from . import backend, io +from vk.utils import io + +from . import backend class Format(Enum): diff --git a/vk/tracking/db/io.py b/vk/tracking/db/io.py deleted file mode 100644 index a89865f..0000000 --- a/vk/tracking/db/io.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) 2017 Egor Tensin <Egor.Tensin@gmail.com> -# This file is part of the "VK scripts" project. -# For details, see https://github.com/egor-tensin/vk-scripts. -# Distributed under the MIT License. - -from contextlib import contextmanager -import csv -import sys - - -class FileWriterCSV: - def __init__(self, fd=sys.stdout): - self._fd = fd - self._writer = csv.writer(fd, lineterminator='\n') - - @staticmethod - def _convert_row_old_python(row): - if isinstance(row, (list, tuple)): - return row - return list(row) - - def write_row(self, row): - if sys.version_info < (3, 5): - row = self._convert_row_old_python(row) - self._writer.writerow(row) - self._fd.flush() - - -class FileReaderCSV: - def __init__(self, fd=sys.stdin): - self._reader = csv.reader(fd) - - def __iter__(self): - return iter(self._reader) - - -@contextmanager -def _open_file(path=None, default=None, **kwargs): - if path is None: - yield default - else: - with open(path, **kwargs) as fd: - yield fd - - -_DEFAULT_ENCODING = 'utf-8' - - -def open_output_text_file(path=None, mode='w'): - return _open_file(path, default=sys.stdout, mode=mode, - encoding=_DEFAULT_ENCODING) - - -def open_input_text_file(path=None): - return _open_file(path, default=sys.stdin, mode='r', - encoding=_DEFAULT_ENCODING) diff --git a/vk/utils/io.py b/vk/utils/io.py index bb8eef9..eccb726 100644 --- a/vk/utils/io.py +++ b/vk/utils/io.py @@ -18,9 +18,19 @@ class FileWriterJSON: self._fd.write('\n') +class FileReaderCSV: + def __init__(self, fd=sys.stdin): + self._reader = csv.reader(fd) + + def __iter__(self): + return iter(self._reader) + + class FileWriterCSV: - def __init__(self, fd=sys.stdout): + def __init__(self, fd=sys.stdout, flush=False): + self._fd = fd self._writer = csv.writer(fd, lineterminator='\n') + self._flush = flush @staticmethod def _convert_row_old_python(row): @@ -32,6 +42,8 @@ class FileWriterCSV: if sys.version_info < (3, 5): row = self._convert_row_old_python(row) self._writer.writerow(row) + if self._flush: + self._fd.flush() @contextmanager @@ -43,8 +55,17 @@ def _open_file(path=None, default=None, **kwargs): yield fd -def open_output_text_file(path=None): - return _open_file(path, default=sys.stdout, mode='w', encoding='utf-8') +_DEFAULT_ENCODING = 'utf-8' + + +def open_input_text_file(path=None): + return _open_file(path, default=sys.stdin, mode='r', + encoding=_DEFAULT_ENCODING) + + +def open_output_text_file(path=None, mode='w'): + return _open_file(path, default=sys.stdout, mode=mode, + encoding=_DEFAULT_ENCODING) def open_output_binary_file(path=None): |