diff options
author | Egor Tensin <Egor.Tensin@gmail.com> | 2016-06-18 02:22:36 +0300 |
---|---|---|
committer | Egor Tensin <Egor.Tensin@gmail.com> | 2016-06-18 02:22:36 +0300 |
commit | f39f18c94060642859900bdb04cb4df2164bc8b4 (patch) | |
tree | eb29d7577e95c32a0894c134214b86db44c3b074 /vk | |
parent | refactoring (diff) | |
download | vk-scripts-f39f18c94060642859900bdb04cb4df2164bc8b4.tar.gz vk-scripts-f39f18c94060642859900bdb04cb4df2164bc8b4.zip |
put format-specific db writers/readers together
Diffstat (limited to 'vk')
-rw-r--r-- | vk/utils/tracking/db/__init__.py | 7 | ||||
-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.py | 63 | ||||
-rw-r--r-- | vk/utils/tracking/db/backend/log.py | 75 | ||||
-rw-r--r-- | vk/utils/tracking/db/backend/null.py | 37 | ||||
-rw-r--r-- | vk/utils/tracking/db/format.py | 35 | ||||
-rw-r--r-- | vk/utils/tracking/db/reader/csv.py | 31 | ||||
-rw-r--r-- | vk/utils/tracking/db/writer/__init__.py | 5 | ||||
-rw-r--r-- | vk/utils/tracking/db/writer/csv.py | 47 | ||||
-rw-r--r-- | vk/utils/tracking/db/writer/log.py | 57 | ||||
-rw-r--r-- | vk/utils/tracking/status_tracker.py | 5 | ||||
-rw-r--r-- | vk/utils/tracking/utils/how_much_online.py | 24 |
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) |