group online periods by date/user/etc.

1 files changed, 46 insertions, 17 deletions

diff --git a/vk/utils/tracking/online_periods.py b/vk/utils/tracking/online_periods.py
index 9a99863..8bd2bda 100644
--- a/vk/utils/tracking/online_periods.py
+++ b/vk/utils/tracking/online_periods.py
@@ -2,45 +2,76 @@
 # This file is licensed under the terms of the MIT License.
 # See LICENSE.txt for details.
 
+from collections import OrderedDict
 from collections.abc import MutableMapping
+from datetime import timedelta
 
 from vk.user import User
 
 class OnlinePeriodEnumerator(MutableMapping):
     def __init__(self):
-        self._records_by_user = {}
+        self._records = {}
 
-    def __getitem__(self, key):
-        return self._records_by_user[self._normalize_key(key)]
+    def __getitem__(self, user):
+        return self._records[user]
 
-    def __setitem__(self, key, value):
-        self._records_by_user[self._normalize_key(key)] = value
+    def __setitem__(self, user, record):
+        self._records[user] = record
 
-    def __delitem__(self, key):
-        del self._records_by_user[self._normalize_key(key)]
+    def __delitem__(self, user):
+        del self._records[user]
 
     def __iter__(self):
-        return iter(self._records_by_user)
+        return iter(self._records)
 
     def __len__(self):
-        return len(self._records_by_user)
-
-    @staticmethod
-    def _normalize_key(key):
-        return key.get_uid() if isinstance(key, User) else key
+        return len(self._records)
 
     def enum(self, db_reader):
         for record in db_reader:
             period = self._insert_record(record)
-            #print(period)
             if period is not None:
                 yield period
 
+    def duration_by_user(self, db_reader):
+        by_user = {}
+        for user, time_from, time_to in self.enum(db_reader):
+            if user not in by_user:
+                by_user[user] = timedelta()
+            by_user[user] += time_to - time_from
+        return by_user
+
+    def duration_by_date(self, db_reader):
+        by_date = OrderedDict()
+        for _, time_from, time_to in self.enum(db_reader):
+            for date, duration in self._enum_dates_and_durations(time_from, time_to):
+                if date not in by_date:
+                    by_date[date] = timedelta()
+                by_date[date] += duration
+        return by_date
+
+    def duration_by_weekday(self, db_reader):
+        by_weekday = OrderedDict()
+        for weekday in range(7):
+            by_weekday[weekday] = timedelta()
+        for _, time_from, time_to in self.enum(db_reader):
+            for date, duration in self._enum_dates_and_durations(time_from, time_to):
+                by_weekday[date.weekday()] += duration
+        return by_weekday
+
+    @staticmethod
+    def _enum_dates_and_durations(time_from, time_to):
+        while time_from.date() != time_to.date():
+            next_day = time_from.date() + timedelta(days=1)
+            yield time_from.date(), next_day - time_from
+            time_from = next_day
+        yield time_to.date(), time_to - time_from
+
     def _insert_record(self, record):
         return self._insert_user(record.to_user())
 
     def _known_user(self, user):
-        return user.get_uid() in self._records_by_user
+        return user.get_uid() in self._records
 
     def _unknown_user(self, user):
         return not self._known_user(user)
@@ -48,10 +79,8 @@ class OnlinePeriodEnumerator(MutableMapping):
     def _insert_user(self, user):
         if user not in self or self[user].is_offline():
             self[user] = user
-            #print(2)
             return None
         if user.is_online():
-            #print(3)
             print(user._fields)
             return None
         period = user, self[user].get_last_seen_time(), user.get_last_seen_time()