diff options
Diffstat (limited to '')
-rw-r--r-- | vk/tracking/db/format.py | 83 |
1 files changed, 30 insertions, 53 deletions
diff --git a/vk/tracking/db/format.py b/vk/tracking/db/format.py index e1b34a1..7b3f312 100644 --- a/vk/tracking/db/format.py +++ b/vk/tracking/db/format.py @@ -3,76 +3,53 @@ # For details, see https://github.com/egor-tensin/vk-scripts. # Distributed under the MIT License. -from contextlib import contextmanager from enum import Enum import sys -from . import backend +from . import backend, io class Format(Enum): CSV = 'csv' LOG = 'log' NULL = 'null' - @contextmanager - def create_writer(self, path=None): - with self._open_output_file(path) as fd: - if self is Format.CSV: - yield backend.csv.Writer(fd) - elif self is Format.LOG: - yield backend.log.Writer(fd) - elif self is Format.NULL: - yield backend.null.Writer() - else: - raise NotImplementedError('unsupported database format: ' + str(self)) + def __str__(self): + return self.value - @contextmanager - def _open_output_file(self, path=None): - fd = sys.stdout - if path is None: - pass - elif self is Format.CSV or self is Format.LOG: - fd = open(path, 'w', encoding='utf-8') + def create_writer(self, fd=sys.stdout): + if self is Format.CSV: + return backend.csv.Writer(fd) + elif self is Format.LOG: + return backend.log.Writer(fd) elif self is Format.NULL: - pass + return backend.null.Writer() else: raise NotImplementedError('unsupported database format: ' + str(self)) - try: - yield fd - finally: - if fd is not sys.stdout: - fd.close() - @contextmanager - def create_reader(self, path=None): - with self._open_input_file(path) as fd: - if self is Format.CSV: - yield backend.csv.Reader(fd) - elif self is Format.LOG: - raise NotImplementedError('cannot read from a log file') - elif self is Format.NULL: - yield backend.null.Reader() - else: - raise NotImplementedError('unsupported database format: ' + str(self)) + def open_output_file(self, path=None): + if self is Format.CSV or self is Format.LOG: + return io.open_output_text_file(path) + elif self is Format.NULL: + return io.open_output_text_file(None) + else: + raise NotImplementedError('unsupported database format: ' + str(self)) - @contextmanager - def _open_input_file(self, path=None): - fd = sys.stdin - if path is None: - pass - elif self is Format.CSV: - fd = open(path, encoding='utf-8') + def create_reader(self, fd=sys.stdin): + if self is Format.CSV: + return backend.csv.Reader(fd) elif self is Format.LOG: - raise NotImplementedError('cannot read from a log file') + return NotImplementedError('cannot read from a log file') elif self is Format.NULL: - pass + return backend.null.Reader() else: raise NotImplementedError('unsupported database format: ' + str(self)) - try: - yield fd - finally: - if fd is not sys.stdin: - fd.close() - def __str__(self): - return self.value + def open_input_file(self, path=None): + if self is Format.CSV: + return io.open_input_text_file(path) + elif self is Format.LOG: + raise NotImplementedError('cannot read from a log file') + elif self is Format.NULL: + return io.open_input_text_file(None) + else: + raise NotImplementedError('unsupported database format: ' + str(self)) |