diff options
Diffstat (limited to '')
| -rw-r--r-- | mutual_friends.py | 19 |
1 files changed, 10 insertions, 9 deletions
diff --git a/mutual_friends.py b/mutual_friends.py index b0bf363..8823619 100644 --- a/mutual_friends.py +++ b/mutual_friends.py @@ -8,7 +8,7 @@ from enum import Enum import json import sys -import vk.api +from vk.api import API, Language from vk.user import UserField OUTPUT_FIELDS = UserField.UID, UserField.FIRST_NAME, UserField.LAST_NAME, UserField.SCREEN_NAME @@ -31,13 +31,13 @@ def print_mutual_friends_csv(mutual_friends): def print_mutual_friends_json(mutual_friends): print(json.dumps([extract_output_fields(user) for user in mutual_friends], indent=3)) -def print_mutual_friends(mutual_friends, output_format): - if output_format is OutputFormat.CSV: +def print_mutual_friends(mutual_friends, fmt): + if fmt is OutputFormat.CSV: print_mutual_friends_csv(mutual_friends) - elif output_format is OutputFormat.JSON: + elif fmt is OutputFormat.JSON: print_mutual_friends_json(mutual_friends) else: - raise NotImplementedError('unsupported output format: ' + str(output_format)) + raise NotImplementedError('unsupported output format: ' + str(fmt)) class OutputFormat(Enum): CSV = 'csv' @@ -60,15 +60,16 @@ if __name__ == '__main__': parser.add_argument(metavar='UID', dest='uids', nargs='+', help='user IDs or "screen names"') - parser.add_argument('-f', '--format', type=output_format, - choices=tuple(x for x in OutputFormat), + parser.add_argument('--output-format', type=output_format, + choices=tuple(fmt for fmt in OutputFormat), default=OutputFormat.CSV, help='specify output format') + args = parser.parse_args() - api = vk.api.API(vk.api.Language.EN) + api = API(Language.EN) users = api.users_get(args.uids) friend_lists = map(lambda user: frozenset(query_friend_list(api, user)), users) mutual_friends = frozenset.intersection(*friend_lists) - print_mutual_friends(mutual_friends, args.format) + print_mutual_friends(mutual_friends, args.output_format) |