From b3c502ffd0a7feb158e9efdbfcf9bd62ce0393b4 Mon Sep 17 00:00:00 2001 From: Egor Tensin Date: Mon, 3 May 2021 21:55:12 +0300 Subject: vk.tracking: merge *_status to status.py --- vk/tracking/status.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) (limited to 'vk/tracking/status.py') diff --git a/vk/tracking/status.py b/vk/tracking/status.py index a16a34e..0a02413 100644 --- a/vk/tracking/status.py +++ b/vk/tracking/status.py @@ -3,13 +3,17 @@ # For details, see https://github.com/egor-tensin/vk-scripts. # Distributed under the MIT License. +import argparse from collections.abc import Callable import contextlib import time import signal +import sys +from vk.api import API import vk.error from vk.user import UserField +from vk.tracking.db import Format as DatabaseFormat class StatusTracker: @@ -137,3 +141,83 @@ class StatusTracker: self._do_loop(uids) except (KeyboardInterrupt, StatusTracker.StopLooping): pass + + +DEFAULT_TIMEOUT = StatusTracker.DEFAULT_TIMEOUT +DEFAULT_DB_FORMAT = DatabaseFormat.CSV + + +def _parse_positive_integer(s): + try: + n = int(s) + except ValueError: + raise argparse.ArgumentTypeError('must be a positive integer: ' + s) + if n < 1: + raise argparse.ArgumentTypeError('must be a positive integer: ' + s) + return n + + +def _parse_database_format(s): + try: + return DatabaseFormat(s) + except ValueError: + raise argparse.ArgumentTypeError('invalid database format: ' + s) + + +def _parse_args(args=None): + if args is None: + args = sys.argv[1:] + + parser = argparse.ArgumentParser( + description='Track when people go online/offline.') + + parser.add_argument('uids', metavar='UID', nargs='+', + help='user IDs or "screen names"') + parser.add_argument('-t', '--timeout', metavar='SECONDS', + type=_parse_positive_integer, + default=DEFAULT_TIMEOUT, + help='set refresh interval') + parser.add_argument('-O', '--only-once', action='store_true', + help='query the status only once and exit') + parser.add_argument('-l', '--log', metavar='PATH', dest='log_path', + help='set log file path (standard output by default)') + parser.add_argument('-f', '--format', dest='db_fmt', + type=_parse_database_format, + choices=DatabaseFormat, + default=DEFAULT_DB_FORMAT, + help='specify database format') + parser.add_argument('-o', '--output', metavar='PATH', dest='db_path', + help='set database file path') + + return parser.parse_args(args) + + +def track_status( + uids, timeout=DEFAULT_TIMEOUT, log_path=None, + db_path=None, db_fmt=DEFAULT_DB_FORMAT, only_once=False): + + api = API() + tracker = StatusTracker(api, timeout) + + if db_fmt is DatabaseFormat.LOG or db_path is None: + db_fmt = DatabaseFormat.NULL + + with DatabaseFormat.LOG.open_output_file(log_path) as log_fd: + log_writer = DatabaseFormat.LOG.create_writer(log_fd) + tracker.add_database_writer(log_writer) + with db_fmt.open_output_file(db_path) as db_fd: + db_writer = db_fmt.create_writer(db_fd) + tracker.add_database_writer(db_writer) + + if only_once: + tracker.query_status(uids) + else: + tracker.loop(uids) + + +def main(args=None): + track_status(**vars(_parse_args(args))) + + +if __name__ == '__main__': + main() -- cgit v1.2.3