aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/vk
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--vk/tracking/db/backend/csv.py14
-rw-r--r--vk/tracking/db/format.py83
-rw-r--r--vk/tracking/db/io.py43
3 files changed, 77 insertions, 63 deletions
diff --git a/vk/tracking/db/backend/csv.py b/vk/tracking/db/backend/csv.py
index e607774..5b2cc7d 100644
--- a/vk/tracking/db/backend/csv.py
+++ b/vk/tracking/db/backend/csv.py
@@ -4,23 +4,20 @@
# Distributed under the MIT License.
from collections.abc import Iterable
-import csv
+from ..io import InputReaderCSV, OutputWriterCSV
from ..record import Record
from ..timestamp import Timestamp
class Writer:
def __init__(self, fd):
- self._fd = fd
- self._writer = csv.writer(fd, lineterminator='\n')
+ self._writer = OutputWriterCSV(fd)
def on_initial_status(self, user):
self._write_record(user)
- self._fd.flush()
def on_status_update(self, user):
self._write_record(user)
- self._fd.flush()
def on_connection_error(self, e):
pass
@@ -28,10 +25,7 @@ class Writer:
def _write_record(self, user):
if not self:
return
- self._write_row(self._record_to_row(Record.from_user(user)))
-
- def _write_row(self, row):
- self._writer.writerow(row)
+ self._writer.write_row(self._record_to_row(Record.from_user(user)))
@staticmethod
def _record_to_row(record):
@@ -39,7 +33,7 @@ class Writer:
class Reader(Iterable):
def __init__(self, fd):
- self._reader = csv.reader(fd)
+ self._reader = InputReaderCSV(fd)
def __iter__(self):
return map(Reader._record_from_row, self._reader)
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))
diff --git a/vk/tracking/db/io.py b/vk/tracking/db/io.py
new file mode 100644
index 0000000..04bab3d
--- /dev/null
+++ b/vk/tracking/db/io.py
@@ -0,0 +1,43 @@
+# 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 OutputWriterCSV:
+ def __init__(self, fd=sys.stdout):
+ self._fd = fd
+ self._writer = csv.writer(fd, lineterminator='\n')
+
+ def write_row(self, row):
+ self._writer.writerow(row)
+ self._fd.flush()
+
+class InputReaderCSV:
+ 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):
+ fd = default
+ if path is None:
+ pass
+ else:
+ fd = open(path, **kwargs)
+ try:
+ yield fd
+ finally:
+ if fd is not default:
+ fd.close()
+
+def open_output_text_file(path=None):
+ return _open_file(path, default=sys.stdout, mode='w', encoding='utf-8')
+
+def open_input_text_file(path=None):
+ return _open_file(path, default=sys.stdin, mode='r', encoding='utf-8')