diff options
Diffstat (limited to '')
| -rw-r--r-- | track_status.py | 292 |
1 files changed, 210 insertions, 82 deletions
diff --git a/track_status.py b/track_status.py index b8b747c..81693bd 100644 --- a/track_status.py +++ b/track_status.py @@ -2,82 +2,201 @@ # This file is licensed under the terms of the MIT License. # See LICENSE.txt for details. -import logging, time +from collections import Callable +import csv +from datetime import datetime +import logging +import sys +import time import vk.api import vk.error -import vk.user - -def format_user(user): - if user.has_last_name(): - return '{} {}'.format(user.get_last_name(), user.get_first_name()) - else: - return '{}'.format(user.get_first_name()) - -def format_user_is_online(user): - return '{} is ONLINE'.format(format_user(user)) - -def format_user_is_offline(user): - return '{} is OFFLINE'.format(format_user(user)) - -def format_user_last_seen(user): - return '{} was last seen at {}'.format(format_user(user), user.get_last_seen()) - -def format_user_went_online(user): - return '{} went ONLINE'.format(format_user(user)) - -def format_user_went_offline(user): - return '{} went OFFLINE'.format(format_user(user)) - -def user_is_online(user): - logging.info(format_user_is_online(user)) - -def user_is_offline(user): - logging.info(format_user_is_offline(user)) - logging.info(format_user_last_seen(user)) - -def user_went_online(user): - logging.info(format_user_went_online(user)) - -def user_went_offline(user): - logging.info(format_user_went_offline(user)) - -def print_status(user): - if user.is_online(): - user_is_online(user) - else: - user_is_offline(user) - -def print_status_update(user): - if user.is_online(): - user_went_online(user) - else: - user_went_offline(user) - -USER_FIELDS = vk.user.Field.ONLINE, vk.user.Field.LAST_SEEN - -def update_status(api, uids): - return {user.get_uid(): user for user in api.users_get(uids, USER_FIELDS)} - -DEFAULT_TIMEOUT=5 +from vk.user import Field + +class CSVWriter: + def __init__(self, path): + if path is None: + self._fd = None + else: + self._fd = open(path, 'w') + self._writer = csv.writer(self._fd, lineterminator='\n') + + def _is_valid(self): + return self._fd is not None + + def __enter__(self): + if not self._is_valid(): + return None + self._fd.__enter__() + return self + + def __exit__(self, *args): + if self._is_valid(): + self._fd.__exit__(*args) + + def flush(self): + if self._is_valid(): + self._fd.flush() + + def write_status(self, user): + self._write_row(self._status_to_row(user)) + self.flush() + + def _write_row(self, row): + self._writer.writerow(row) + + @staticmethod + def _status_to_row(user): + return [ + datetime.utcnow().replace(microsecond=0).isoformat(), + user.get_uid(), + user.get_first_name(), + user.get_last_name(), + user.get_screen_name(), + user.is_online(), + ] + +class Logger: + @staticmethod + def set_up(stream=sys.stdout): + logging.basicConfig(format='[%(asctime)s] %(message)s', + stream=stream, + level=logging.INFO, + datefmt='%Y-%m-%d %H:%M:%S') + + @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)) + + @staticmethod + def on_exception(e): + logging.exception(e) -def loop_update_status(api, uids, timeout=DEFAULT_TIMEOUT): - users = update_status(api, uids) - for uid in users: - print_status(users[uid]) - while True: - time.sleep(timeout) + @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 {}'.format(Logger._format_user(user), user.get_last_seen()) + + @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)) + +class StatusTracker: + DEFAULT_TIMEOUT = 5 + + def __init__(self, api, timeout=DEFAULT_TIMEOUT): + self._api = api + self._timeout = timeout + self._on_initial_status = [] + self._on_status_update = [] + self._on_connection_error = [] + + def _wait_after_connection_error(self): + time.sleep(self._timeout) + + def add_initial_status_handler(self, fn): + self._assert_is_callback(fn) + self._on_initial_status.append(fn) + + def add_status_update_handler(self, fn): + self._assert_is_callback(fn) + self._on_status_update.append(fn) + + def add_connection_error_handler(self, fn): + self._assert_is_callback(fn) + self._on_connection_error.append(fn) + + @staticmethod + def _assert_is_callback(fn): + if not isinstance(fn, Callable): + raise TypeError() + + USER_FIELDS = Field.SCREEN_NAME, Field.ONLINE, Field.LAST_SEEN + + def _query_status(self, uids): + return {user.get_uid(): user for user in self._api.users_get(uids, StatusTracker.USER_FIELDS)} + + def _notify_status(self, user): + for fn in self._on_initial_status: + fn(user) + + def _notify_status_update(self, user): + for fn in self._on_status_update: + fn(user) + + def _notify_connection_error(self, e): + for fn in self._on_connection_error: + fn(e) + + def _query_initial_status(self, uids): + while True: + try: + return self._query_status(uids) + except vk.error.ConnectionError as e: + self._notify_connection_error(e) + self._wait_after_connection_error() + + def _query_status_updates(self, uids): + while True: + self._wait_after_connection_error() + try: + return self._query_status(uids) + except vk.error.ConnectionError as e: + self._notify_connection_error(e) + + @staticmethod + def _filter_status_updates(old_users, new_users): + for uid, user in new_users.items(): + if old_users[uid].is_online() != user.is_online(): + old_users[uid] = user + yield user + + def _do_loop(self, uids): + users = self._query_initial_status(uids) + for user in users.values(): + self._notify_status(user) + while True: + updated_users = self._query_status_updates(uids) + for user in self._filter_status_updates(users, updated_users): + self._notify_status_update(user) + + def loop(self, uids): try: - updated_users = update_status(api, uids) - except vk.error.ConnectionError: - continue - for uid in updated_users: - if users[uid].is_online() != updated_users[uid].is_online(): - users[uid] = updated_users[uid] - print_status_update(updated_users[uid]) + self._do_loop(uids) + except KeyboardInterrupt: + pass if __name__ == '__main__': - import argparse, sys + import argparse def natural_number(s): x = int(s) @@ -90,25 +209,34 @@ if __name__ == '__main__': parser.add_argument(metavar='UID', dest='uids', nargs='+', help='user IDs or "screen names"') - parser.add_argument('-t', '--timeout', default=DEFAULT_TIMEOUT, - type=natural_number, + parser.add_argument('-t', '--timeout', type=natural_number, + default=StatusTracker.DEFAULT_TIMEOUT, help='set refresh interval (seconds)') parser.add_argument('-l', '--log', type=argparse.FileType('w'), default=sys.stdout, help='set log file path (stdout by default)') + parser.add_argument('-o', '--output', default=None, + help='set status database path') + args = parser.parse_args() - logging.basicConfig(format='[%(asctime)s] %(message)s', - stream=args.log, - level=logging.INFO, - datefmt='%Y-%m-%d %H:%M:%S') + Logger.set_up(args.log) api = vk.api.API(vk.api.Language.EN) + tracker = StatusTracker(api, args.timeout) - try: - loop_update_status(api, args.uids, timeout=args.timeout) - except KeyboardInterrupt: - pass - except Exception as e: - logging.exception(e) - sys.exit(1) + tracker.add_initial_status_handler(Logger.on_initial_status) + tracker.add_status_update_handler(Logger.on_status_update) + tracker.add_connection_error_handler(Logger.on_exception) + + with CSVWriter(args.output) as csv_writer: + + if csv_writer is not None: + tracker.add_initial_status_handler(lambda user: csv_writer.write_status(user)) + tracker.add_status_update_handler(lambda user: csv_writer.write_status(user)) + + try: + tracker.loop(args.uids) + except Exception as e: + Logger.on_exception(e) + sys.exit(1) |