summaryrefslogtreecommitdiff
path: root/wikibania/ban/BanDBWrapper.py
blob: 8f396b07b5c15ba89ba1244443740d935b96d68f (plain)
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
from collections import defaultdict

import numpy


class BanDBWrapper:
    def __init__(self, ban_db):
        self.ban_db = ban_db

    def get_all_durations(self):
        return [ban.calc_duration() for ban in self.ban_db.list()]

    def get_all_countries(self):
        return [ban.lookup_country_code() for ban in self.ban_db.list()]

    def get_durations_by_country(self):
        return [(ban.lookup_country_code(), ban.calc_duration()) for ban in self.ban_db.list()]

    def calc_average_duration_by_country(self):
        ban_durations_by_country = defaultdict(list)

        for country, ban_duration in self.get_durations_by_country():
            ban_durations_by_country[country].append(ban_duration)

        return {country: numpy.mean(ban_durations) for country, ban_durations in ban_durations_by_country.items()}