From eb930123454771b80465505579d723c92b3dd84c Mon Sep 17 00:00:00 2001 From: Egor Tensin Date: Thu, 16 Jun 2016 23:32:24 +0300 Subject: refactoring & support more user fields And a bunch of other minor improvements. --- vk/utils/tracking/db/reader/csv.py | 14 ++++- vk/utils/tracking/db/record.py | 114 +++++++++++++++++++++++++------------ vk/utils/tracking/db/writer/csv.py | 6 +- 3 files changed, 93 insertions(+), 41 deletions(-) (limited to 'vk/utils/tracking/db') diff --git a/vk/utils/tracking/db/reader/csv.py b/vk/utils/tracking/db/reader/csv.py index b66e397..e9c9407 100644 --- a/vk/utils/tracking/db/reader/csv.py +++ b/vk/utils/tracking/db/reader/csv.py @@ -2,11 +2,12 @@ # This file is licensed under the terms of the MIT License. # See LICENSE.txt for details. +from collections.abc import Iterable import csv -from ..record import Record +from ..record import Record, Timestamp -class Reader: +class Reader(Iterable): def __init__(self, path): self._fd = open(path) self._reader = csv.reader(self._fd) @@ -19,4 +20,11 @@ class Reader: self._fd.__exit__(*args) def __iter__(self): - return map(Record.from_row, self._reader) + return map(Reader._record_from_row, self._reader) + + @staticmethod + def _record_from_row(row): + record = Record(Timestamp.from_string(row[0])) + for i in range(len(Record.FIELDS)): + record[Record.FIELDS[i]] = row[i + 1] + return record diff --git a/vk/utils/tracking/db/record.py b/vk/utils/tracking/db/record.py index 7cb054f..4748a37 100644 --- a/vk/utils/tracking/db/record.py +++ b/vk/utils/tracking/db/record.py @@ -3,15 +3,51 @@ # See LICENSE.txt for details. from collections import OrderedDict +from collections.abc import MutableMapping from datetime import datetime, timezone -from vk.user import Field as UserField +from vk.user import LastSeen, User, UserField -def _gen_timestamp(): - return datetime.now(timezone.utc).replace(microsecond=0) +class Timestamp: + @staticmethod + def _new(): + return datetime.utcnow() + + @staticmethod + def _is_timezone_aware(dt): + return dt.tzinfo is not None and dt.tzinfo.utcoffset(dt) is not None -class Record: - _USER_FIELDS = ( + @staticmethod + def _lose_timezone(dt): + if Timestamp._is_timezone_aware(dt): + return dt.astimezone(timezone.utc).replace(tzinfo=None) + return dt + + def __init__(self, dt=None): + if dt is None: + dt = self._new() + dt = dt.replace(microsecond=0) + dt = self._lose_timezone(dt) + self._dt = dt + + @staticmethod + def from_string(s): + return Timestamp(datetime.strptime(s, '%Y-%m-%dT%H:%M:%SZ')) + + def __str__(self): + return self._dt.isoformat() + 'Z' + + @staticmethod + def from_last_seen(ls): + return Timestamp(ls.get_time()) + + def to_last_seen(self): + ls = LastSeen() + ls.set_time(self._dt) + return ls + +class Record(MutableMapping): + FIELDS = ( UserField.UID, UserField.FIRST_NAME, UserField.LAST_NAME, @@ -20,48 +56,52 @@ class Record: UserField.LAST_SEEN, ) - def __init__(self, fields, timestamp=None): + def __init__(self, timestamp=None, fields=None): + if timestamp is None: + timestamp = Timestamp() + if fields is None: + fields = OrderedDict() + self._timestamp = timestamp self._fields = fields - self._timestamp = timestamp if timestamp is not None else _gen_timestamp() - - def __iter__(self): - return iter(self._fields) - - def __contains__(self, field): - return field in self._fields def __getitem__(self, field): + if field is UserField.LAST_SEEN: + return Timestamp.from_last_seen(self._fields[field]) return self._fields[field] def __setitem__(self, field, value): - self._fields[field] = value + if field is UserField.LAST_SEEN: + if isinstance(value, str): + value = Timestamp.from_string(value).to_last_seen() + elif isinstance(value, Timestamp): + value = value.to_last_seen() + elif isinstance(value, LastSeen): + pass + else: + raise TypeError() + self._fields[field] = User.parse(field, value) - def get_timestamp(self): - return self._timestamp + def __delitem__(self, field): + del self._fields[field] - @staticmethod - def _timestamp_from_string(s): - return datetime.fromtimestamp(s) + def __iter__(self): + return iter(self._fields) + + def __len__(self): + return len(self._fields) - def timestamp_to_string(self): - return self.get_timestamp().isoformat() + def get_timestamp(self): + return self._timestamp @staticmethod def from_user(user): - fields = OrderedDict() - for field in Record._USER_FIELDS: - fields[field] = user[field] - if UserField.LAST_SEEN in Record._USER_FIELDS: - fields[UserField.LAST_SEEN] = fields[UserField.LAST_SEEN].isoformat() - return Record(fields) + instance = Record() + for field in Record.FIELDS: + instance[field] = user[field] + return instance - @staticmethod - def from_row(row): - timestamp = Record._timestamp_from_string(row[0]) - fields = OrderedDict() - for i in range(len(Record._USER_FIELDS)): - fields[Record._USER_FIELDS[i]] = row[i + 1] - return Record(fields, timestamp) - - def to_row(self): - return [self.timestamp_to_string()] + [self[field] for field in self] + def to_user(self): + user = User() + for field in self: + user[field] = self[field] + return user diff --git a/vk/utils/tracking/db/writer/csv.py b/vk/utils/tracking/db/writer/csv.py index 8dc2299..8c635b4 100644 --- a/vk/utils/tracking/db/writer/csv.py +++ b/vk/utils/tracking/db/writer/csv.py @@ -36,8 +36,12 @@ class Writer: def write_record(self, user): if not self: return - self._write_row(Record.from_user(user).to_row()) + self._write_row(self._record_to_row(Record.from_user(user))) self.flush() def _write_row(self, row): self._writer.writerow(row) + + @staticmethod + def _record_to_row(record): + return [str(record.get_timestamp())] + [str(record[field]) for field in record] -- cgit v1.2.3