put format-specific db writers/readers together

12 files changed, 228 insertions, 160 deletions

diff --git a/vk/utils/tracking/db/__init__.py b/vk/utils/tracking/db/__init__.py
index e69de29..9e6b74a 100644
--- a/vk/utils/tracking/db/__init__.py
+++ b/vk/utils/tracking/db/__init__.py
@@ -0,0 +1,7 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from .format import Format
+
+__all__ = 'format'
diff --git a/vk/utils/tracking/db/reader/__init__.py b/vk/utils/tracking/db/backend/__init__.py
index 330d5a8..4b3c278 100644
--- a/vk/utils/tracking/db/reader/__init__.py
+++ b/vk/utils/tracking/db/backend/__init__.py
@@ -2,4 +2,4 @@
 # This file is licensed under the terms of the MIT License.
 # See LICENSE.txt for details.
 
-__all__ = 'csv',
+__all__ = 'csv', 'log', 'null'
diff --git a/vk/utils/tracking/db/backend/csv.py b/vk/utils/tracking/db/backend/csv.py
new file mode 100644
index 0000000..10a504f
--- /dev/null
+++ b/vk/utils/tracking/db/backend/csv.py
@@ -0,0 +1,63 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# 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 ..timestamp import Timestamp
+
+class Writer:
+    def __init__(self, fd):
+        self._fd = fd
+        self._writer = csv.writer(fd, lineterminator='\n')
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def on_initial_status(self, user):
+        self._write_record(user)
+        self._fd.flush()
+
+    def on_status_update(self, user):
+        self._write_record(user)
+        self._fd.flush()
+
+    def on_connection_error(self, e):
+        pass
+
+    def _write_record(self, user):
+        if not self:
+            return
+        self._write_row(self._record_to_row(Record.from_user(user)))
+
+    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]
+
+class Reader(Iterable):
+    def __init__(self, fd):
+        self._reader = csv.reader(fd)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def __iter__(self):
+        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/backend/log.py b/vk/utils/tracking/db/backend/log.py
new file mode 100644
index 0000000..625257b
--- /dev/null
+++ b/vk/utils/tracking/db/backend/log.py
@@ -0,0 +1,75 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+import logging
+
+class Writer:
+    def __init__(self, fd):
+        self._logger = logging.getLogger(__file__)
+        self._logger.setLevel(logging.INFO)
+        handler = logging.StreamHandler(fd)
+        handler.setFormatter(logging.Formatter(
+            fmt='[%(asctime)s] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S'))
+        self._logger.addHandler(handler)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def info(self, msg):
+        self._logger.info(msg)
+
+    def exception(self, e):
+        self._logger.exception(e)
+
+    def on_initial_status(self, user):
+        if user.is_online():
+            self.info(self._format_user_is_online(user))
+        else:
+            self.info(self._format_user_is_offline(user))
+        self.info(self._format_user_last_seen(user))
+
+    def on_status_update(self, user):
+        if user.is_online():
+            self.info(self._format_user_went_online(user))
+        else:
+            self.info(self._format_user_went_offline(user))
+        self.info(self._format_user_last_seen(user))
+
+    def on_connection_error(self, e):
+        #self.exception(e)
+        pass
+
+    @staticmethod
+    def _format_user(user):
+        if user.has_last_name():
+            return '{} {}'.format(user.get_first_name(), user.get_last_name())
+        else:
+            return '{}'.format(user.get_first_name())
+
+    @staticmethod
+    def _format_user_is_online(user):
+        return '{} is ONLINE.'.format(Writer._format_user(user))
+
+    @staticmethod
+    def _format_user_is_offline(user):
+        return '{} is OFFLINE.'.format(Writer._format_user(user))
+
+    @staticmethod
+    def _format_user_last_seen(user):
+        return '{} was last seen at {} using {}.'.format(
+            Writer._format_user(user),
+            user.get_last_seen_time_local(),
+            user.get_last_seen_platform().get_description_for_sentence())
+
+    @staticmethod
+    def _format_user_went_online(user):
+        return '{} went ONLINE.'.format(Writer._format_user(user))
+
+    @staticmethod
+    def _format_user_went_offline(user):
+        return '{} went OFFLINE.'.format(Writer._format_user(user))
diff --git a/vk/utils/tracking/db/backend/null.py b/vk/utils/tracking/db/backend/null.py
new file mode 100644
index 0000000..139a9f0
--- /dev/null
+++ b/vk/utils/tracking/db/backend/null.py
@@ -0,0 +1,37 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from collections.abc import Iterable
+
+class Writer:
+    def __init__(self, fd):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def on_initial_status(self, user):
+        pass
+
+    def on_status_update(self, user):
+        pass
+
+    def on_connection_error(self, e):
+        pass
+
+class Reader(Iterable):
+    def __init__(self, fd):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def __iter__(self):
+        pass
diff --git a/vk/utils/tracking/db/format.py b/vk/utils/tracking/db/format.py
new file mode 100644
index 0000000..4856094
--- /dev/null
+++ b/vk/utils/tracking/db/format.py
@@ -0,0 +1,35 @@
+# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
+# This file is licensed under the terms of the MIT License.
+# See LICENSE.txt for details.
+
+from enum import Enum
+
+from .backend import *
+
+class Format(Enum):
+    CSV = 'csv'
+    LOG = 'log'
+    NULL = 'null'
+
+    def create_writer(self, fd):
+        if self is Format.CSV:
+            return csv.Writer(fd)
+        elif self is Format.LOG:
+            return log.Writer(fd)
+        elif self is Format.NULL:
+            return null.Writer(fd)
+        else:
+            raise NotImplementedError('unsupported database format: ' + str(self))
+
+    def create_reader(self, fd):
+        if self is Format.CSV:
+            return csv.Reader(fd)
+        elif self is Format.LOG:
+            raise NotImplementedError()
+        elif self is Format.NULL:
+            return null.Reader(fd)
+        else:
+            raise NotImplementedError('unsupported database format: ' + str(self))
+
+    def __str__(self):
+        return self.value
diff --git a/vk/utils/tracking/db/reader/csv.py b/vk/utils/tracking/db/reader/csv.py
deleted file mode 100644
index 9d4e7dc..0000000
--- a/vk/utils/tracking/db/reader/csv.py
+++ /dev/null
@@ -1,31 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# 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 ..timestamp import Timestamp
-
-class Reader(Iterable):
-    def __init__(self, path):
-        self._fd = open(path)
-        self._reader = csv.reader(self._fd)
-
-    def __enter__(self):
-        self._fd.__enter__()
-        return self
-
-    def __exit__(self, *args):
-        self._fd.__exit__(*args)
-
-    def __iter__(self):
-        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/writer/__init__.py b/vk/utils/tracking/db/writer/__init__.py
deleted file mode 100644
index 8eb62e1..0000000
--- a/vk/utils/tracking/db/writer/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-__all__ = 'csv', 'log'
diff --git a/vk/utils/tracking/db/writer/csv.py b/vk/utils/tracking/db/writer/csv.py
deleted file mode 100644
index 8c635b4..0000000
--- a/vk/utils/tracking/db/writer/csv.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-import csv
-
-from ..record import Record
-
-class Writer:
-    def __init__(self, path, mode='w'):
-        if path is None:
-            self._fd = None
-        else:
-            self._fd = open(path, mode)
-            self._writer = csv.writer(self._fd, lineterminator='\n')
-
-    def __bool__(self):
-        return self._fd is not None
-
-    def __enter__(self):
-        if not self:
-            return self
-        self._fd.__enter__()
-        return self
-
-    def __exit__(self, *args):
-        if not self:
-            return
-        self._fd.__exit__(*args)
-
-    def flush(self):
-        if not self:
-            return
-        self._fd.flush()
-
-    def write_record(self, user):
-        if not self:
-            return
-        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]
diff --git a/vk/utils/tracking/db/writer/log.py b/vk/utils/tracking/db/writer/log.py
deleted file mode 100644
index faef29b..0000000
--- a/vk/utils/tracking/db/writer/log.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# Copyright 2016 Egor Tensin <Egor.Tensin@gmail.com>
-# This file is licensed under the terms of the MIT License.
-# See LICENSE.txt for details.
-
-import logging
-
-class Logger:
-    @staticmethod
-    def on_initial_status(user):
-        if user.is_online():
-            logging.info(Logger._format_user_is_online(user))
-        else:
-            logging.info(Logger._format_user_is_offline(user))
-        logging.info(Logger._format_user_last_seen(user))
-
-    @staticmethod
-    def on_status_update(user):
-        if user.is_online():
-            logging.info(Logger._format_user_went_online(user))
-        else:
-            logging.info(Logger._format_user_went_offline(user))
-        logging.info(Logger._format_user_last_seen(user))
-
-    @staticmethod
-    def on_connection_error(e):
-        #logging.exception(e)
-        pass
-
-    @staticmethod
-    def _format_user(user):
-        if user.has_last_name():
-            return '{} {}'.format(user.get_first_name(), user.get_last_name())
-        else:
-            return '{}'.format(user.get_first_name())
-
-    @staticmethod
-    def _format_user_is_online(user):
-        return '{} is ONLINE.'.format(Logger._format_user(user))
-
-    @staticmethod
-    def _format_user_is_offline(user):
-        return '{} is OFFLINE.'.format(Logger._format_user(user))
-
-    @staticmethod
-    def _format_user_last_seen(user):
-        return '{} was last seen at {} using {}.'.format(
-            Logger._format_user(user),
-            user.get_last_seen_time_local(),
-            user.get_last_seen_platform().get_description_for_sentence())
-
-    @staticmethod
-    def _format_user_went_online(user):
-        return '{} went ONLINE.'.format(Logger._format_user(user))
-
-    @staticmethod
-    def _format_user_went_offline(user):
-        return '{} went OFFLINE.'.format(Logger._format_user(user))
diff --git a/vk/utils/tracking/status_tracker.py b/vk/utils/tracking/status_tracker.py
index f208884..3d1f032 100644
--- a/vk/utils/tracking/status_tracker.py
+++ b/vk/utils/tracking/status_tracker.py
@@ -21,6 +21,11 @@ class StatusTracker:
     def _wait_after_connection_error(self):
         time.sleep(self._timeout)
 
+    def add_database_writer(self, writer):
+        self.add_initial_status_handler(lambda user: writer.on_initial_status(user))
+        self.add_status_update_handler(lambda user: writer.on_status_update(user))
+        self.add_connection_error_handler(lambda e: writer.on_connection_error(e))
+
     def add_initial_status_handler(self, fn):
         self._assert_is_callback(fn)
         self._on_initial_status.append(fn)
diff --git a/vk/utils/tracking/utils/how_much_online.py b/vk/utils/tracking/utils/how_much_online.py
index bf7fa25..b40c357 100644
--- a/vk/utils/tracking/utils/how_much_online.py
+++ b/vk/utils/tracking/utils/how_much_online.py
@@ -12,6 +12,7 @@ import sys
 import matplotlib.pyplot as plt
 import numpy as np
 
+from ..db import Format as DatabaseFormat
 from vk.user import UserField
 
 def process_database(db_reader, writer):
@@ -25,22 +26,6 @@ def process_database(db_reader, writer):
     for user, wasted_time in wasted_time_by_user.items():
         writer.write_wasted_time(user, wasted_time)
 
-class DatabaseFormat(Enum):
-    CSV = 'csv'
-
-    def __str__(self):
-        return self.value
-
-def open_database_csv(path):
-    from vk.utils.tracking.db.reader.csv import Reader
-    return Reader(path)
-
-def open_database(path, fmt):
-    if fmt is DatabaseFormat.CSV:
-        return open_database_csv(path)
-    else:
-        raise NotImplementedError('unsupported database format: ' + str(fmt))
-
 class OutputFormat(Enum):
     CSV = 'csv'
     JSON = 'json'
@@ -277,12 +262,13 @@ if __name__ == '__main__':
         except ValueError:
             raise argparse.ArgumentTypeError()
 
-    parser.add_argument('input', help='database path')
+    parser.add_argument('input', type=argparse.FileType('r'),
+                        help='database path')
     parser.add_argument('output', type=argparse.FileType('w'),
                         nargs='?', default=sys.stdout,
                         help='output path (standard output by default)')
     parser.add_argument('--input-format', type=database_format,
-                        choices=tuple(str(fmt) for fmt in DatabaseFormat),
+                        choices=tuple(fmt for fmt in DatabaseFormat),
                         default=DatabaseFormat.CSV,
                         help='specify database format')
     parser.add_argument('--output-format', type=output_format,
@@ -292,6 +278,6 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
 
-    with open_database(args.input, args.input_format) as db_reader:
+    with args.input_format.create_reader(args.input) as db_reader:
         with open_output_writer(args.output, args.output_format) as output_writer:
             process_database(db_reader, output_writer)