diff options
Diffstat (limited to '')
| -rw-r--r-- | bin/mutual_friends.py | 58 | ||||
| -rw-r--r-- | docs/mutual_friends.md | 18 |
2 files changed, 49 insertions, 27 deletions
diff --git a/bin/mutual_friends.py b/bin/mutual_friends.py index 8823619..b72b317 100644 --- a/bin/mutual_friends.py +++ b/bin/mutual_friends.py @@ -11,7 +11,7 @@ import sys from vk.api import API, Language from vk.user import UserField -OUTPUT_FIELDS = UserField.UID, UserField.FIRST_NAME, UserField.LAST_NAME, UserField.SCREEN_NAME +OUTPUT_FIELDS = UserField.UID, UserField.FIRST_NAME, UserField.LAST_NAME def query_friend_list(api, user): return api.friends_get(user.get_uid(), fields=OUTPUT_FIELDS) @@ -22,22 +22,35 @@ def extract_output_fields(user): new_user[str(field)] = user[field] if field in user else None return new_user -def print_mutual_friends_csv(mutual_friends): - writer = csv.writer(sys.stdout, lineterminator='\n') - for user in mutual_friends: - user = extract_output_fields(user) - writer.writerow(user.values()) +class OutputWriterCSV: + def __init__(self, fd=sys.stdout): + self._writer = csv.writer(fd, lineterminator='\n') -def print_mutual_friends_json(mutual_friends): - print(json.dumps([extract_output_fields(user) for user in mutual_friends], indent=3)) + def __enter__(self): + return self -def print_mutual_friends(mutual_friends, fmt): - if fmt is OutputFormat.CSV: - print_mutual_friends_csv(mutual_friends) - elif fmt is OutputFormat.JSON: - print_mutual_friends_json(mutual_friends) - else: - raise NotImplementedError('unsupported output format: ' + str(fmt)) + def __exit__(self, *args): + pass + + def write_mutual_friends(self, friend_list): + for user in mutual_friends: + user = extract_output_fields(user) + self._writer.writerow(user.values()) + +class OutputWriterJSON: + def __init__(self, fd=sys.stdout): + self._fd = fd + self._array = [] + + def __enter__(self): + return self + + def __exit__(self, *args): + self._fd.write(json.dumps(self._array, indent=3, ensure_ascii=False)) + + def write_mutual_friends(self, friend_list): + for user in friend_list: + self._array.append(extract_output_fields(user)) class OutputFormat(Enum): CSV = 'csv' @@ -46,6 +59,14 @@ class OutputFormat(Enum): def __str__(self): return self.value + def create_writer(self, fd=sys.stdout): + if self is OutputFormat.CSV: + return OutputWriterCSV(fd) + elif self is OutputFormat.JSON: + return OutputWriterJSON(fd) + else: + raise NotImplementedError('unsupported output format: ' + str(self)) + if __name__ == '__main__': import argparse @@ -64,6 +85,9 @@ if __name__ == '__main__': choices=tuple(fmt for fmt in OutputFormat), default=OutputFormat.CSV, help='specify output format') + parser.add_argument('--output', type=argparse.FileType('w', encoding='utf-8'), + default=sys.stdout, + help='set output file path (standard output by default)') args = parser.parse_args() @@ -72,4 +96,6 @@ if __name__ == '__main__': friend_lists = map(lambda user: frozenset(query_friend_list(api, user)), users) mutual_friends = frozenset.intersection(*friend_lists) - print_mutual_friends(mutual_friends, args.output_format) + + with args.output_format.create_writer(args.output) as writer: + writer.write_mutual_friends(mutual_friends) diff --git a/docs/mutual_friends.md b/docs/mutual_friends.md index cc9396e..4e0e4b4 100644 --- a/docs/mutual_friends.md +++ b/docs/mutual_friends.md @@ -6,23 +6,21 @@ Learn who your ex and her new boyfriend are both friends with. Usage ----- -Run from the top-level directory using `python -m`. -For example: +Run from the top-level directory using `python -m`: > python -m bin.mutual_friends -h - usage: mutual_friends.py [-h] [--output-format {csv,json}] UID [UID ...] + usage: mutual_friends.py [-h] [--output-format {csv,json}] [--output OUTPUT] + UID [UID ...] ... For example (using made up user IDs/"screen names"), > python -m bin.mutual_friends john.doe jane.doe - 89497105,John,Smith,john.smith - 3698577,Jane,Smith,jane.smith + 89497105,John,Smith + 3698577,Jane,Smith In the example above, both "John Doe" and "Jane Doe" are friends with "John Smith" and "Jane Smith", whose user IDs are 89497105 and 3698577 respectively. -Their "screen names" (the part after "vk.com/" of their personal page URLs) are -"john.smith" and "jane.smith". The output format is CSV (comma-separated values) by default. You can also get a JSON document: @@ -32,14 +30,12 @@ You can also get a JSON document: { "uid": 89497105, "first_name": "John", - "last_name": "Smith", - "screen_name": "john.smith" + "last_name": "Smith" }, { "uid": 3698577, "first_name": "Jane", - "last_name": "Smith", - "screen_name": "jane.smith" + "last_name": "Smith" } ] |