From 9f81ed27c6f84d7d0fef1638577091e87f0f4864 Mon Sep 17 00:00:00 2001
From: M1ha <work_shvein_mihail@mail.ru>
Date: Tue, 31 Jan 2017 18:13:46 +0500
Subject: [PATCH] 1) Added to_dict model method 2) Fixed bug in test_freeze
 cleaning, if backups don't contain all directory names (e. g. 1, 2, 3, 6, 7 -
 count=5, created_backup=8, not 6)

---
 src/infi/clickhouse_orm/models.py | 12 ++++++++++++
 tests/test_models.py              | 24 ++++++++++++++++++++++++
 tests/test_system_models.py       | 12 ++++++------
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/infi/clickhouse_orm/models.py b/src/infi/clickhouse_orm/models.py
index 09127e3..78033d4 100644
--- a/src/infi/clickhouse_orm/models.py
+++ b/src/infi/clickhouse_orm/models.py
@@ -161,3 +161,15 @@ class Model(with_metaclass(ModelBase)):
         fields = [f for f in self._fields if f[1].is_insertable()] if insertable_only else self._fields
         return '\t'.join(field.to_db_string(data[name], quote=False) for name, field in fields)
 
+    def to_dict(self, insertable_only=False, field_names=None):
+        '''
+        Returns the instance's column values as a dict.
+        :param bool insertable_only: If True, returns only fields, that can be inserted into database
+        :param field_names: An iterable of field names to return
+        '''
+        fields = [f for f in self._fields if f[1].is_insertable()] if insertable_only else self._fields
+        if field_names is not None:
+            fields = [f for f in fields if f[0] in field_names]
+
+        data = self.__dict__
+        return {name: field.to_python(data[name]) for name, field in fields}
diff --git a/tests/test_models.py b/tests/test_models.py
index 4b259a9..239f0c8 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -55,6 +55,29 @@ class ModelTestCase(unittest.TestCase):
         instance.int_field = '99'
         self.assertEquals(instance.int_field, 99)
 
+    def test_to_dict(self):
+        instance = SimpleModel(date_field='1973-12-06', int_field='100', float_field='7')
+        self.assertDictEqual(instance.to_dict(), {
+            "date_field": datetime.date(1973, 12, 6),
+            "int_field": 100,
+            "float_field": 7.0,
+            "datetime_field": datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+            "alias_field": 0.0,
+            'str_field': 'dozo'
+        })
+        self.assertDictEqual(instance.to_dict(insertable_only=True), {
+            "date_field": datetime.date(1973, 12, 6),
+            "int_field": 100,
+            "float_field": 7.0,
+            "datetime_field": datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc),
+            'str_field': 'dozo'
+        })
+        self.assertDictEqual(
+            instance.to_dict(insertable_only=True, field_names=('int_field', 'alias_field', 'datetime_field')), {
+                "int_field": 100,
+                "datetime_field": datetime.datetime(1970, 1, 1, 0, 0, 0, tzinfo=pytz.utc)
+            })
+
 
 class SimpleModel(Model):
 
@@ -63,6 +86,7 @@ class SimpleModel(Model):
     str_field = StringField(default='dozo')
     int_field = Int32Field(default=17)
     float_field = Float32Field()
+    alias_field = Float32Field(alias='float_field')
 
     engine = MergeTree('date_field', ('int_field', 'date_field'))
 
diff --git a/tests/test_system_models.py b/tests/test_system_models.py
index db97e5d..d8e075a 100644
--- a/tests/test_system_models.py
+++ b/tests/test_system_models.py
@@ -21,9 +21,9 @@ class SystemPartTest(unittest.TestCase):
     def tearDown(self):
         self.database.drop_database()
 
-    def _get_backups_count(self):
+    def _get_backups(self):
         _, dirnames, _ = next(os.walk(self.BACKUP_DIR))
-        return len(dirnames)
+        return dirnames
 
     def test_get_all(self):
         parts = SystemPart.all(self.database)
@@ -51,12 +51,12 @@ class SystemPartTest(unittest.TestCase):
     def test_freeze(self):
         parts = list(SystemPart.all(self.database))
         # There can be other backups in the folder
-        backups_count = self._get_backups_count()
+        prev_backups = set(self._get_backups())
         parts[0].freeze(self.database)
-        backup_number = self._get_backups_count()
-        self.assertEqual(backup_number, backups_count + 1)
+        backups = set(self._get_backups())
+        self.assertEqual(len(backups), len(prev_backups) + 1)
         # Clean created backup
-        shutil.rmtree(self.BACKUP_DIR + '{0}'.format(backup_number))
+        shutil.rmtree(self.BACKUP_DIR + '{0}'.format(list(backups - prev_backups)[0]))
 
     def test_fetch(self):
         # TODO Not tested, as I have no replication set