diff options
Diffstat (limited to 'bin/mutual_friends.py')
| -rw-r--r-- | bin/mutual_friends.py | 17 |
1 files changed, 14 insertions, 3 deletions
diff --git a/bin/mutual_friends.py b/bin/mutual_friends.py index 82015bb..550d957 100644 --- a/bin/mutual_friends.py +++ b/bin/mutual_friends.py @@ -14,22 +14,27 @@ 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) @@ -38,6 +43,7 @@ class OutputSinkCSV(OutputSinkMutualFriends): 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) @@ -45,6 +51,7 @@ class OutputSinkJSON(OutputSinkMutualFriends): def write_mutual_friends(self, friend_list): self._writer.write(friend_list) + class OutputFormat(Enum): CSV = 'csv' JSON = 'json' @@ -59,10 +66,10 @@ class OutputFormat(Enum): def create_sink(self, fd=sys.stdout): if self is OutputFormat.CSV: return OutputSinkCSV(fd) - elif self is OutputFormat.JSON: + if self is OutputFormat.JSON: return OutputSinkJSON(fd) - else: - raise NotImplementedError('unsupported output format: ' + str(self)) + raise NotImplementedError('unsupported output format: ' + str(self)) + def _parse_output_format(s): try: @@ -70,6 +77,7 @@ def _parse_output_format(s): except ValueError: raise argparse.ArgumentTypeError('invalid output format: ' + s) + def _parse_args(args=None): if args is None: args = sys.argv[1:] @@ -89,6 +97,7 @@ def _parse_args(args=None): return parser.parse_args(args) + def write_mutual_friends(uids, out_path=None, out_fmt=OutputFormat.CSV): api = API() users = api.users_get(uids) @@ -101,8 +110,10 @@ def write_mutual_friends(uids, out_path=None, out_fmt=OutputFormat.CSV): 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() |