aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/bin/mutual_friends.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--bin/mutual_friends.py17
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()