aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/mutual_friends.py
diff options
context:
space:
mode:
authorEgor Tensin <Egor.Tensin@gmail.com>2016-06-15 20:06:12 +0300
committerEgor Tensin <Egor.Tensin@gmail.com>2016-06-15 20:06:12 +0300
commit2031ea3ca49c5834d3ee6fb252369802ed482a6b (patch)
tree1b6e8082dda3485ec7badb1e5c033a9589753a01 /mutual_friends.py
parentmove API wrappers to a package (diff)
downloadvk-scripts-2031ea3ca49c5834d3ee6fb252369802ed482a6b.tar.gz
vk-scripts-2031ea3ca49c5834d3ee6fb252369802ed482a6b.zip
mutual_friends.py: switch to CSV/JSON output
Diffstat (limited to '')
-rw-r--r--mutual_friends.py64
1 files changed, 52 insertions, 12 deletions
diff --git a/mutual_friends.py b/mutual_friends.py
index 2f32038..3b86a8e 100644
--- a/mutual_friends.py
+++ b/mutual_friends.py
@@ -2,33 +2,73 @@
# This file is licensed under the terms of the MIT License.
# See LICENSE.txt for details.
+from collections import OrderedDict
+import csv
+from enum import Enum
+import json
+import sys
+
import vk.api
-import vk.user
+from vk.user import Field
+
+def query_friend_list(api, user):
+ return api.friends_get(user.get_uid(), fields=Field.SCREEN_NAME)
+
+OUTPUT_FIELDS = Field.FIRST_NAME, Field.LAST_NAME, Field.SCREEN_NAME
+
+def extract_output_fields(user):
+ new_user = OrderedDict()
+ for field in OUTPUT_FIELDS:
+ new_user[str(field)] = user[field] if field in user else None
+ return new_user
-def query_friends(api, user):
- return api.friends_get(user.get_uid(), fields=vk.user.Field.SCREEN_NAME)
+def print_mutual_friends_csv(mutual_friends):
+ writer = csv.writer(sys.stdout)
+ for user in mutual_friends:
+ user = extract_output_fields(user)
+ writer.writerow(user.values())
-def format_user(user):
- if user.has_last_name():
- return '{} {} ({})'.format(user.get_last_name(), user.get_first_name(), user.get_screen_name())
+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:
+ print_mutual_friends_csv(mutual_friends)
+ elif output_format is OutputFormat.JSON:
+ print_mutual_friends_json(mutual_friends)
else:
- return '{} ({})'.format(user.get_first_name(), user.get_screen_name())
+ raise NotImplementedError('unsupported output format: ' + str(output_format))
+
+class OutputFormat(Enum):
+ CSV = 'csv'
+ JSON = 'json'
+
+ def __str__(self):
+ return self.value
if __name__ == '__main__':
import argparse
+ def output_format(s):
+ try:
+ return OutputFormat(s)
+ except ValueError:
+ raise argparse.ArgumentError()
+
parser = argparse.ArgumentParser(
description='Learn who your ex and her new boyfriend are both friends with.')
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),
+ default=OutputFormat.CSV,
+ help='specify output format')
args = parser.parse_args()
api = vk.api.API(vk.api.Language.EN)
- users = api.users_get(args.uids, fields=vk.user.Field.SCREEN_NAME)
+ users = api.users_get(args.uids, fields=Field.SCREEN_NAME)
- friend_lists = map(lambda user: frozenset(query_friends(api, user)), users)
+ friend_lists = map(lambda user: frozenset(query_friend_list(api, user)), users)
mutual_friends = frozenset.intersection(*friend_lists)
- if mutual_friends:
- for friend in mutual_friends:
- print(format_user(friend))
+ print_mutual_friends(mutual_friends, args.format)