1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
|
# Copyright (c) 2015 Egor Tensin <Egor.Tensin@gmail.com>
# This file is part of the "VK scripts" project.
# For details, see https://github.com/egor-tensin/vk-scripts.
# Distributed under the MIT License.
import abc
import argparse
from collections import OrderedDict
from enum import Enum
import sys
from vk.api import API
from vk.user import UserField
from .utils import io
_OUTPUT_USER_FIELDS = UserField.UID, UserField.FIRST_NAME, UserField.LAST_NAME
def _query_friend_list(api, user):
return api.friends_get(user.get_uid(), fields=_OUTPUT_USER_FIELDS)
def _filter_user_fields(user):
new_user = OrderedDict()
for field in _OUTPUT_USER_FIELDS:
new_user[str(field)] = user[field] if field in user else None
return new_user
class OutputSinkMutualFriends(metaclass=abc.ABCMeta):
@abc.abstractmethod
def write_mutual_friends(self, friend_list):
pass
class OutputSinkCSV(OutputSinkMutualFriends):
def __init__(self, fd=sys.stdout):
self._writer = io.FileWriterCSV(fd)
def write_mutual_friends(self, friend_list):
for user in friend_list:
self._writer.write_row(user.values())
class OutputSinkJSON(OutputSinkMutualFriends):
def __init__(self, fd=sys.stdout):
self._writer = io.FileWriterJSON(fd)
def write_mutual_friends(self, friend_list):
self._writer.write(friend_list)
class OutputFormat(Enum):
CSV = 'csv'
JSON = 'json'
def __str__(self):
return self.value
@staticmethod
def open_file(path=None):
return io.open_output_text_file(path)
def create_sink(self, fd=sys.stdout):
if self is OutputFormat.CSV:
return OutputSinkCSV(fd)
elif self is OutputFormat.JSON:
return OutputSinkJSON(fd)
else:
raise NotImplementedError('unsupported output format: ' + str(self))
def _parse_output_format(s):
try:
return OutputFormat(s)
except ValueError:
raise argparse.ArgumentTypeError('invalid output format: ' + s)
def _parse_args(args=None):
if args is None:
args = sys.argv[1:]
parser = argparse.ArgumentParser(
description='Learn who your ex and her new boyfriend are both friends with.')
parser.add_argument('uids', metavar='UID', nargs='+',
help='user IDs or "screen names"')
parser.add_argument('-f', '--format', dest='out_fmt',
type=_parse_output_format,
default=OutputFormat.CSV,
choices=OutputFormat,
help='specify output format')
parser.add_argument('-o', '--output', metavar='PATH', dest='out_path',
help='set output file path (standard output by default)')
return parser.parse_args(args)
def write_mutual_friends(uids, out_path=None, out_fmt=OutputFormat.CSV):
api = API()
users = api.users_get(uids)
friend_lists = (frozenset(_query_friend_list(api, user)) for user in users)
mutual_friends = frozenset.intersection(*friend_lists)
mutual_friends = [_filter_user_fields(user) for user in mutual_friends]
with out_fmt.open_file(out_path) as out_fd:
sink = out_fmt.create_sink(out_fd)
sink.write_mutual_friends(mutual_friends)
def main(args=None):
write_mutual_friends(**vars(_parse_args(args)))
if __name__ == '__main__':
main()
|
