From f7bdf2547bd8aad9113b014331882e228d3365f9 Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Thu, 27 Mar 2014 13:39:35 +0200
Subject: [PATCH] Add tests for WeekDaysType

---
 .travis.yml                        |  1 +
 docs/data_types.rst                |  9 ++++++
 sqlalchemy_utils/__init__.py       |  2 ++
 sqlalchemy_utils/types/__init__.py |  4 +--
 sqlalchemy_utils/types/bit.py      | 20 +++++++++++++
 sqlalchemy_utils/types/weekdays.py | 12 ++++++--
 tests/types/test_weekdays.py       | 45 ++++++++++++++++++++++++++++++
 7 files changed, 87 insertions(+), 6 deletions(-)
 create mode 100644 sqlalchemy_utils/types/bit.py
 create mode 100644 tests/types/test_weekdays.py

diff --git a/.travis.yml b/.travis.yml
index 764c980..51e8a52 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -3,6 +3,7 @@ addons:
 
 before_script:
   - psql -c 'create database sqlalchemy_utils_test;' -U postgres
+  - mysql -e 'create database sqlalchemy_continuum_test;'
 
 
 language: python
diff --git a/docs/data_types.rst b/docs/data_types.rst
index dac0a53..2bee9f4 100644
--- a/docs/data_types.rst
+++ b/docs/data_types.rst
@@ -96,3 +96,12 @@ UUIDType
 
 .. autoclass:: UUIDType
 
+
+
+WeekDaysType
+^^^^^^^^^^^^
+
+.. module:: sqlalchemy_utils.types.weekdays
+
+.. autoclass:: WeekDaysType
+
diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py
index a90c194..ea504bc 100644
--- a/sqlalchemy_utils/__init__.py
+++ b/sqlalchemy_utils/__init__.py
@@ -53,6 +53,7 @@ from .types import (
     TSVectorType,
     URLType,
     UUIDType,
+    WeekDaysType
 )
 
 
@@ -113,4 +114,5 @@ __all__ = (
     TSVectorType,
     URLType,
     UUIDType,
+    WeekDaysType
 )
diff --git a/sqlalchemy_utils/types/__init__.py b/sqlalchemy_utils/types/__init__.py
index 14344db..471121f 100644
--- a/sqlalchemy_utils/types/__init__.py
+++ b/sqlalchemy_utils/types/__init__.py
@@ -21,7 +21,7 @@ from .timezone import TimezoneType
 from .ts_vector import TSVectorType
 from .url import URLType
 from .uuid import UUIDType
-from .weekdays import WeekDay, WeekDays, WeekDaysType
+from .weekdays import WeekDaysType
 
 
 __all__ = (
@@ -49,8 +49,6 @@ __all__ = (
     TSVectorType,
     URLType,
     UUIDType,
-    WeekDay,
-    WeekDays,
     WeekDaysType,
 )
 
diff --git a/sqlalchemy_utils/types/bit.py b/sqlalchemy_utils/types/bit.py
new file mode 100644
index 0000000..6cc6942
--- /dev/null
+++ b/sqlalchemy_utils/types/bit.py
@@ -0,0 +1,20 @@
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import BIT
+
+
+class BitType(sa.types.TypeDecorator):
+    """
+    BitType offers way of saving BITs into database.
+    """
+    impl = sa.types.BINARY
+
+    def __init__(self, length=1, **kwargs):
+        self.length = length
+        sa.types.TypeDecorator.__init__(self, **kwargs)
+
+    def load_dialect_impl(self, dialect):
+        # Use the native BIT type for drivers that has it.
+        if dialect.name == 'postgresql':
+            return dialect.type_descriptor(BIT(self.length))
+        else:
+            return dialect.type_descriptor(type(self.impl)(self.length))
diff --git a/sqlalchemy_utils/types/weekdays.py b/sqlalchemy_utils/types/weekdays.py
index 505988c..4be83ec 100644
--- a/sqlalchemy_utils/types/weekdays.py
+++ b/sqlalchemy_utils/types/weekdays.py
@@ -1,11 +1,12 @@
 import six
 from sqlalchemy import types
-from sqlalchemy.dialects.postgresql import BIT
 from sqlalchemy_utils.primitives import WeekDay, WeekDays
+from .scalar_coercible import ScalarCoercible
+from .bit import BitType
 
 
-class WeekDaysType(types.TypeDecorator):
-    impl = BIT(WeekDay.NUM_WEEK_DAYS)
+class WeekDaysType(types.TypeDecorator, ScalarCoercible):
+    impl = BitType(WeekDay.NUM_WEEK_DAYS)
 
     def process_bind_param(self, value, dialect):
         if isinstance(value, WeekDays):
@@ -17,3 +18,8 @@ class WeekDaysType(types.TypeDecorator):
     def process_result_value(self, value, dialect):
         if value is not None:
             return WeekDays(value)
+
+    def _coerce(self, value):
+        if value is not None and not isinstance(value, WeekDays):
+            return WeekDays(value)
+        return value
diff --git a/tests/types/test_weekdays.py b/tests/types/test_weekdays.py
new file mode 100644
index 0000000..cfc3266
--- /dev/null
+++ b/tests/types/test_weekdays.py
@@ -0,0 +1,45 @@
+import sqlalchemy as sa
+from sqlalchemy_utils.types import WeekDaysType
+from sqlalchemy_utils.primitives import WeekDays
+
+from tests import TestCase
+
+
+class WeekDaysTypeTestCase(TestCase):
+    def create_models(self):
+        class Schedule(self.Base):
+            __tablename__ = 'schedule'
+            id = sa.Column(sa.Integer, primary_key=True)
+            working_days = sa.Column(WeekDaysType)
+
+            def __repr__(self):
+                return 'Schedule(%r)' % self.id
+
+        self.Schedule = Schedule
+
+    def test_color_parameter_processing(self):
+        schedule = self.Schedule(
+            working_days='0001111'
+        )
+        self.session.add(schedule)
+        self.session.commit()
+
+        schedule = self.session.query(self.Schedule).first()
+        assert isinstance(schedule.working_days, WeekDays)
+
+    def test_scalar_attributes_get_coerced_to_objects(self):
+        schedule = self.Schedule(working_days=u'1010101')
+
+        assert isinstance(schedule.working_days, WeekDays)
+
+
+class TestWeekDaysTypeOnSQLite(WeekDaysTypeTestCase):
+    dns = 'sqlite:///:memory:'
+
+
+class TestWeekDaysTypeOnPostgres(WeekDaysTypeTestCase):
+    dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
+
+
+class TestWeekDaysTypeOnMySQL(WeekDaysTypeTestCase):
+    dns = 'mysql+pymysql://travis@localhost/sqlalchemy_utils_test'