aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/vk/utils
diff options
context:
space:
mode:
authorEgor Tensin <Egor.Tensin@gmail.com>2016-06-18 02:22:36 +0300
committerEgor Tensin <Egor.Tensin@gmail.com>2016-06-18 02:22:36 +0300
commitf39f18c94060642859900bdb04cb4df2164bc8b4 (patch)
treeeb29d7577e95c32a0894c134214b86db44c3b074 /vk/utils
parentrefactoring (diff)
downloadvk-scripts-f39f18c94060642859900bdb04cb4df2164bc8b4.tar.gz
vk-scripts-f39f18c94060642859900bdb04cb4df2164bc8b4.zip
put format-specific db writers/readers together
Diffstat (limited to '')
-rw-r--r--vk/utils/tracking/db/__init__.py7
-rw-r--r--vk/utils/tracking/db/backend/__init__.py (renamed from vk/utils/tracking/db/reader/__init__.py)2
-rw-r--r--vk/utils/tracking/db/backend/csv.py63
-rw-r--r--vk/utils/tracking/db/backend/log.py75
-rw-r--r--vk/utils/tracking/db/backend/null.py37
-rw-r--r--vk/utils/tracking/db/format.py35
-rw-r--r--vk/utils/tracking/db/reader/csv.py31
-rw-r--r--vk/utils/tracking/db/writer/__init__.py5
-rw-r--r--vk/utils/tracking/db/writer/csv.py47
-rw-r--r--vk/utils/tracking/db/writer/log.py57
-rw-r--r--vk/utils/tracking/status_tracker.py5
-rw-r--r--vk/utils/tracking/utils/how_much_online.py24
12 files changed, 228 insertions, 160 deletions
diff --git a/vk/utils/tracking/db/__init__.py b/vk/utils/tracking/db/__init__.py
index e69de29..9e6b74a 100644
--- a/vk/utils/tracking/db/__init__.py
+++ b/vk/utils/tracking/db/__init__.py
@@ -0,0 +1,7 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from .format import Format
+
+__all__ = 'format'
diff --git a/vk/utils/tracking/db/reader/__init__.py b/vk/utils/tracking/db/backend/__init__.py
index 330d5a8..4b3c278 100644
--- a/vk/utils/tracking/db/reader/__init__.py
+++ b/vk/utils/tracking/db/backend/__init__.py
@@ -2,4 +2,4 @@
# This file is licensed under the terms of the MIT License.
# See LICENSE.txt for details.
-__all__ = 'csv',
+__all__ = 'csv', 'log', 'null'
diff --git a/vk/utils/tracking/db/backend/csv.py b/vk/utils/tracking/db/backend/csv.py
new file mode 100644
index 0000000..10a504f
--- /dev/null
+++ b/vk/utils/tracking/db/backend/csv.py
@@ -0,0 +1,63 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from collections.abc import Iterable
+import csv
+
+from ..record import Record
+from ..timestamp import Timestamp
+
+class Writer:
+ def __init__(self, fd):
+ self._fd = fd
+ self._writer = csv.writer(fd, lineterminator='\n')
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ pass
+
+ 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
+
+ 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)
+
+ @staticmethod
+ def _record_to_row(record):
+ return [str(record.get_timestamp())] + [str(record[field]) for field in record]
+
+class Reader(Iterable):
+ def __init__(self, fd):
+ self._reader = csv.reader(fd)
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ pass
+
+ def __iter__(self):
+ return map(Reader._record_from_row, self._reader)
+
+ @staticmethod
+ def _record_from_row(row):
+ record = Record(Timestamp.from_string(row[0]))
+ for i in range(len(Record.FIELDS)):
+ record[Record.FIELDS[i]] = row[i + 1]
+ return record
diff --git a/vk/utils/tracking/db/backend/log.py b/vk/utils/tracking/db/backend/log.py
new file mode 100644
index 0000000..625257b
--- /dev/null
+++ b/vk/utils/tracking/db/backend/log.py
@@ -0,0 +1,75 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+import logging
+
+class Writer:
+ def __init__(self, fd):
+ self._logger = logging.getLogger(__file__)
+ self._logger.setLevel(logging.INFO)
+ handler = logging.StreamHandler(fd)
+ handler.setFormatter(logging.Formatter(
+ fmt='[%(asctime)s] %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S'))
+ self._logger.addHandler(handler)
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ pass
+
+ def info(self, msg):
+ self._logger.info(msg)
+
+ def exception(self, e):
+ self._logger.exception(e)
+
+ def on_initial_status(self, user):
+ if user.is_online():
+ self.info(self._format_user_is_online(user))
+ else:
+ self.info(self._format_user_is_offline(user))
+ self.info(self._format_user_last_seen(user))
+
+ def on_status_update(self, user):
+ if user.is_online():
+ self.info(self._format_user_went_online(user))
+ else:
+ self.info(self._format_user_went_offline(user))
+ self.info(self._format_user_last_seen(user))
+
+ def on_connection_error(self, e):
+ #self.exception(e)
+ pass
+
+ @staticmethod
+ def _format_user(user):
+ if user.has_last_name():
+ return '{} {}'.format(user.get_first_name(), user.get_last_name())
+ else:
+ return '{}'.format(user.get_first_name())
+
+ @staticmethod
+ def _format_user_is_online(user):
+ return '{} is ONLINE.'.format(Writer._format_user(user))
+
+ @staticmethod
+ def _format_user_is_offline(user):
+ return '{} is OFFLINE.'.format(Writer._format_user(user))
+
+ @staticmethod
+ def _format_user_last_seen(user):
+ return '{} was last seen at {} using {}.'.format(
+ Writer._format_user(user),
+ user.get_last_seen_time_local(),
+ user.get_last_seen_platform().get_description_for_sentence())
+
+ @staticmethod
+ def _format_user_went_online(user):
+ return '{} went ONLINE.'.format(Writer._format_user(user))
+
+ @staticmethod
+ def _format_user_went_offline(user):
+ return '{} went OFFLINE.'.format(Writer._format_user(user))
diff --git a/vk/utils/tracking/db/backend/null.py b/vk/utils/tracking/db/backend/null.py
new file mode 100644
index 0000000..139a9f0
--- /dev/null
+++ b/vk/utils/tracking/db/backend/null.py
@@ -0,0 +1,37 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from collections.abc import Iterable
+
+class Writer:
+ def __init__(self, fd):
+ pass
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ pass
+
+ def on_initial_status(self, user):
+ pass
+
+ def on_status_update(self, user):
+ pass
+
+ def on_connection_error(self, e):
+ pass
+
+class Reader(Iterable):
+ def __init__(self, fd):
+ pass
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ pass
+
+ def __iter__(self):
+ pass
diff --git a/vk/utils/tracking/db/format.py b/vk/utils/tracking/db/format.py
new file mode 100644
index 0000000..4856094
--- /dev/null
+++ b/vk/utils/tracking/db/format.py
@@ -0,0 +1,35 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from enum import Enum
+
+from .backend import *
+
+class Format(Enum):
+ CSV = 'csv'
+ LOG = 'log'
+ NULL = 'null'
+
+ def create_writer(self, fd):
+ if self is Format.CSV:
+ return csv.Writer(fd)
+ elif self is Format.LOG:
+ return log.Writer(fd)
+ elif self is Format.NULL:
+ return null.Writer(fd)
+ else:
+ raise NotImplementedError('unsupported database format: ' + str(self))
+
+ def create_reader(self, fd):
+ if self is Format.CSV:
+ return csv.Reader(fd)
+ elif self is Format.LOG:
+ raise NotImplementedError()
+ elif self is Format.NULL:
+ return null.Reader(fd)
+ else:
+ raise NotImplementedError('unsupported database format: ' + str(self))
+
+ def __str__(self):
+ return self.value
diff --git a/vk/utils/tracking/db/reader/csv.py b/vk/utils/tracking/db/reader/csv.py
deleted file mode 100644
index 9d4e7dc..0000000
--- a/vk/utils/tracking/db/reader/csv.py
+++ /dev/null
@@ -1,31 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-from collections.abc import Iterable
-import csv
-
-from ..record import Record
-from ..timestamp import Timestamp
-
-class Reader(Iterable):
- def __init__(self, path):
- self._fd = open(path)
- self._reader = csv.reader(self._fd)
-
- def __enter__(self):
- self._fd.__enter__()
- return self
-
- def __exit__(self, *args):
- self._fd.__exit__(*args)
-
- def __iter__(self):
- return map(Reader._record_from_row, self._reader)
-
- @staticmethod
- def _record_from_row(row):
- record = Record(Timestamp.from_string(row[0]))
- for i in range(len(Record.FIELDS)):
- record[Record.FIELDS[i]] = row[i + 1]
- return record
diff --git a/vk/utils/tracking/db/writer/__init__.py b/vk/utils/tracking/db/writer/__init__.py
deleted file mode 100644
index 8eb62e1..0000000
--- a/vk/utils/tracking/db/writer/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-__all__ = 'csv', 'log'
diff --git a/vk/utils/tracking/db/writer/csv.py b/vk/utils/tracking/db/writer/csv.py
deleted file mode 100644
index 8c635b4..0000000
--- a/vk/utils/tracking/db/writer/csv.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-import csv
-
-from ..record import Record
-
-class Writer:
- def __init__(self, path, mode='w'):
- if path is None:
- self._fd = None
- else:
- self._fd = open(path, mode)
- self._writer = csv.writer(self._fd, lineterminator='\n')
-
- def __bool__(self):
- return self._fd is not None
-
- def __enter__(self):
- if not self:
- return self
- self._fd.__enter__()
- return self
-
- def __exit__(self, *args):
- if not self:
- return
- self._fd.__exit__(*args)
-
- def flush(self):
- if not self:
- return
- self._fd.flush()
-
- def write_record(self, user):
- if not self:
- return
- self._write_row(self._record_to_row(Record.from_user(user)))
- self.flush()
-
- def _write_row(self, row):
- self._writer.writerow(row)
-
- @staticmethod
- def _record_to_row(record):
- return [str(record.get_timestamp())] + [str(record[field]) for field in record]
diff --git a/vk/utils/tracking/db/writer/log.py b/vk/utils/tracking/db/writer/log.py
deleted file mode 100644
index faef29b..0000000
--- a/vk/utils/tracking/db/writer/log.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-import logging
-
-class Logger:
- @staticmethod
- def on_initial_status(user):
- if user.is_online():
- logging.info(Logger._format_user_is_online(user))
- else:
- logging.info(Logger._format_user_is_offline(user))
- logging.info(Logger._format_user_last_seen(user))
-
- @staticmethod
- def on_status_update(user):
- if user.is_online():
- logging.info(Logger._format_user_went_online(user))
- else:
- logging.info(Logger._format_user_went_offline(user))
- logging.info(Logger._format_user_last_seen(user))
-
- @staticmethod
- def on_connection_error(e):
- #logging.exception(e)
- pass
-
- @staticmethod
- def _format_user(user):
- if user.has_last_name():
- return '{} {}'.format(user.get_first_name(), user.get_last_name())
- else:
- return '{}'.format(user.get_first_name())
-
- @staticmethod
- def _format_user_is_online(user):
- return '{} is ONLINE.'.format(Logger._format_user(user))
-
- @staticmethod
- def _format_user_is_offline(user):
- return '{} is OFFLINE.'.format(Logger._format_user(user))
-
- @staticmethod
- def _format_user_last_seen(user):
- return '{} was last seen at {} using {}.'.format(
- Logger._format_user(user),
- user.get_last_seen_time_local(),
- user.get_last_seen_platform().get_description_for_sentence())
-
- @staticmethod
- def _format_user_went_online(user):
- return '{} went ONLINE.'.format(Logger._format_user(user))
-
- @staticmethod
- def _format_user_went_offline(user):
- return '{} went OFFLINE.'.format(Logger._format_user(user))
diff --git a/vk/utils/tracking/status_tracker.py b/vk/utils/tracking/status_tracker.py
index f208884..3d1f032 100644
--- a/vk/utils/tracking/status_tracker.py
+++ b/vk/utils/tracking/status_tracker.py
@@ -21,6 +21,11 @@ class StatusTracker:
def _wait_after_connection_error(self):
time.sleep(self._timeout)
+ def add_database_writer(self, writer):
+ self.add_initial_status_handler(lambda user: writer.on_initial_status(user))
+ self.add_status_update_handler(lambda user: writer.on_status_update(user))
+ self.add_connection_error_handler(lambda e: writer.on_connection_error(e))
+
def add_initial_status_handler(self, fn):
self._assert_is_callback(fn)
self._on_initial_status.append(fn)
diff --git a/vk/utils/tracking/utils/how_much_online.py b/vk/utils/tracking/utils/how_much_online.py
index bf7fa25..b40c357 100644
--- a/vk/utils/tracking/utils/how_much_online.py
+++ b/vk/utils/tracking/utils/how_much_online.py
@@ -12,6 +12,7 @@ import sys
import matplotlib.pyplot as plt
import numpy as np
+from ..db import Format as DatabaseFormat
from vk.user import UserField
def process_database(db_reader, writer):
@@ -25,22 +26,6 @@ def process_database(db_reader, writer):
for user, wasted_time in wasted_time_by_user.items():
writer.write_wasted_time(user, wasted_time)
-class DatabaseFormat(Enum):
- CSV = 'csv'
-
- def __str__(self):
- return self.value
-
-def open_database_csv(path):
- from vk.utils.tracking.db.reader.csv import Reader
- return Reader(path)
-
-def open_database(path, fmt):
- if fmt is DatabaseFormat.CSV:
- return open_database_csv(path)
- else:
- raise NotImplementedError('unsupported database format: ' + str(fmt))
-
class OutputFormat(Enum):
CSV = 'csv'
JSON = 'json'
@@ -277,12 +262,13 @@ if __name__ == '__main__':
except ValueError:
raise argparse.ArgumentTypeError()
- parser.add_argument('input', help='database path')
+ parser.add_argument('input', type=argparse.FileType('r'),
+ help='database path')
parser.add_argument('output', type=argparse.FileType('w'),
nargs='?', default=sys.stdout,
help='output path (standard output by default)')
parser.add_argument('--input-format', type=database_format,
- choices=tuple(str(fmt) for fmt in DatabaseFormat),
+ choices=tuple(fmt for fmt in DatabaseFormat),
default=DatabaseFormat.CSV,
help='specify database format')
parser.add_argument('--output-format', type=output_format,
@@ -292,6 +278,6 @@ if __name__ == '__main__':
args = parser.parse_args()
- with open_database(args.input, args.input_format) as db_reader:
+ with args.input_format.create_reader(args.input) as db_reader:
with open_output_writer(args.output, args.output_format) as output_writer:
process_database(db_reader, output_writer)