1 files changed, 25 insertions, 0 deletions

diff --git a/wikibania/ban/BanDBWrapper.py b/wikibania/ban/BanDBWrapper.py
new file mode 100644
index 0000000..8f396b0
--- /dev/null
+++ b/wikibania/ban/BanDBWrapper.py
@@ -0,0 +1,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()}