added Date and Time column types

Author: dantownsend

piccolo/columns/column_types.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 import copy
-from datetime import datetime
+from datetime import datetime, date, time
 import decimal
 import typing as t
 import uuid
 
 from piccolo.columns.base import Column, OnDelete, OnUpdate, ForeignKeyMeta
 from piccolo.columns.operators.string import ConcatPostgres, ConcatSQLite
-from piccolo.custom_types import Datetime, UUIDDefault
+from piccolo.custom_types import UUIDDefault
 from piccolo.querystring import Unquoted, QueryString
 
 if t.TYPE_CHECKING:
-    from piccolo.table import Table  # noqa
-    from piccolo.custom_types import Datetime, UUID_
+    from piccolo.table import Table
+    from piccolo.custom_types import DatetimeArg, UUIDArg, TimeArg, DateArg
 
 
 ###############################################################################
@@ -195,20 +195,11 @@ class UUID(Column):
 
     value_type = uuid.UUID
 
-    def __init__(self, default: UUID_ = UUIDDefault.uuid4, **kwargs) -> None:
+    def __init__(self, default: UUIDArg = UUIDDefault.uuid4, **kwargs) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
 
-    @property
-    def column_type(self):
-        engine_type = self._meta.table._meta.db.engine_type
-        if engine_type == "postgres":
-            return "UUID"
-        elif engine_type == "sqlite":
-            return "VARCHAR"
-        raise Exception("Unrecognized engine type")
-
 
 class Integer(Column):
 
@@ -353,16 +344,40 @@ def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
 
 
+###############################################################################
+
+
 class Timestamp(Column):
 
     value_type = datetime
 
-    def __init__(self, default: Datetime = None, **kwargs) -> None:
+    def __init__(self, default: DatetimeArg = None, **kwargs) -> None:
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+
+class Date(Column):
+    value_type = date
+
+    def __init__(self, default: DatetimeArg = None, **kwargs) -> None:
         self.default = default
         kwargs.update({"default": default})
         super().__init__(**kwargs)
 
 
+class Time(Column):
+    value_type = time
+
+    def __init__(self, default: TimeArg = None, **kwargs) -> None:
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+
+###############################################################################
+
+
 class Boolean(Column):
 
     value_type = bool

piccolo/custom_types.py

@@ -13,6 +13,9 @@
 Iterable = t.Iterable[t.Any]
 
 
+###############################################################################
+
+
 @dataclass
 class Default:
     python: t.Callable
@@ -28,22 +31,29 @@ class DatetimeDefault(Enum):
     )
 
 
+class DateDefault(Enum):
+    now = Default(
+        python=lambda: datetime.datetime.now().date(),
+        postgres="current_date",
+        sqlite="current_date",
+    )
+
+
+class TimeDefault(Enum):
+    now = Default(
+        python=lambda: datetime.datetime.now().time(),
+        postgres="current_time",
+        sqlite="current_time",
+    )
+
+
 class UUIDDefault(Enum):
     uuid4 = Default(
         python=uuid.uuid4, postgres="uuid_generate_v4()", sqlite="''"
     )
 
 
-UUID_ = t.Union[UUIDDefault, uuid.UUID]
-Datetime = t.Union[datetime.datetime, DatetimeDefault, None]
-
-
-__all__ = (
-    "Combinable",
-    "Iterable",
-    "Datetime",
-    "DatetimeDefault",
-    "Default",
-    "UUID_",
-    "UUIDDefault",
-)
+UUIDArg = t.Union[UUIDDefault, uuid.UUID]
+DatetimeArg = t.Union[datetime.datetime, DatetimeDefault, None]
+DateArg = t.Union[datetime.date, DateDefault, None]
+TimeArg = t.Union[datetime.time, TimeDefault, None]

piccolo/engine/postgres.py

@@ -253,7 +253,7 @@ def prep_database(self):
             print(
                 "Unable to create uuid-ossp extension - UUID columns might "
                 "not behave as expected. Make sure your database user has "
-                "permission to create extensions."
+                "permission to create extensions, or add it manually."
             )
 
     ###########################################################################

piccolo/engine/sqlite.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 import contextvars
 from dataclasses import dataclass
+import datetime
 from decimal import Decimal
 import os
 import sqlite3
@@ -33,49 +34,65 @@ def convert_uuid_in(value):
     """
     Converts the UUID value being passed into sqlite.
     """
-    return "uuid:" + str(value)
+    return str(value)
 
 
-def convert_numeric_out(value: bytes):
+def convert_time_in(value):
+    """
+    Converts the time value being passed into sqlite.
+    """
+    return value.isoformat()
+
+
+def convert_date_in(value):
+    """
+    Converts the date value being passed into sqlite.
+    """
+    return value.isoformat()
+
+
+def convert_numeric_out(value: bytes) -> Decimal:
     """
     Convert float values into Decimals.
     """
     return Decimal(value.decode("ascii"))
 
 
-def convert_int_out(value: bytes):
+def convert_int_out(value: bytes) -> int:
     """
     Make sure Integer values are actually of type int.
     """
     return int(float(value))
 
 
-def convert_uuid_out(value: bytes):
+def convert_uuid_out(value: bytes) -> uuid.UUID:
     """
     If the value is a uuid, convert it to a UUID instance.
+    """
+    return uuid.UUID(value.decode("utf8"))
+
+
+def convert_date_out(value: bytes) -> datetime.date:
+    return datetime.date.fromisoformat(value.decode("utf8"))
+
 
-    Performance wise, this isn't great, but SQLite isn't the preferred solution
-    in production, so it's acceptable.
+def convert_time_out(value: bytes) -> datetime.time:
     """
-    decoded = value.decode("utf8")
-    if decoded.startswith("uuid:"):
-        uuid_string = decoded.split("uuid:", 1)[1]
-        try:
-            _uuid = uuid.UUID(uuid_string)
-        except ValueError:
-            return decoded
-        else:
-            return _uuid
-    else:
-        return decoded
+    If the value is a uuid, convert it to a UUID instance.
+    """
+    return datetime.time.fromisoformat(value.decode("utf8"))
 
 
 sqlite3.register_converter("Numeric", convert_numeric_out)
 sqlite3.register_converter("Integer", convert_int_out)
-sqlite3.register_converter("Varchar", convert_uuid_out)
+sqlite3.register_converter("UUID", convert_uuid_out)
+sqlite3.register_converter("Date", convert_date_out)
+sqlite3.register_converter("Time", convert_time_out)
 
 sqlite3.register_adapter(Decimal, convert_numeric_in)
 sqlite3.register_adapter(uuid.UUID, convert_uuid_in)
+sqlite3.register_adapter(datetime.time, convert_time_in)
+sqlite3.register_adapter(datetime.date, convert_date_in)
 
 ###############################################################################
 

tests/columns/test_date.py

@@ -0,0 +1,25 @@
+import datetime
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Date
+
+
+class MyTable(Table):
+    created_on = Date()
+
+
+class TestDate(TestCase):
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_timestamp(self):
+        created_on = datetime.datetime.now().date()
+        row = MyTable(created_on=created_on)
+        row.save().run_sync()
+
+        result = MyTable.objects().first().run_sync()
+        self.assertEqual(result.created_on, created_on)

tests/columns/test_time.py

@@ -0,0 +1,25 @@
+import datetime
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Time
+
+
+class MyTable(Table):
+    created_on = Time()
+
+
+class TestTime(TestCase):
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_timestamp(self):
+        created_on = datetime.datetime.now().time()
+        row = MyTable(created_on=created_on)
+        row.save().run_sync()
+
+        result = MyTable.objects().first().run_sync()
+        self.assertEqual(result.created_on, created_on)

tests/columns/test_timestamp.py

@@ -0,0 +1,25 @@
+import datetime
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Timestamp
+
+
+class MyTable(Table):
+    created_on = Timestamp()
+
+
+class TestTimestamp(TestCase):
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_timestamp(self):
+        created_on = datetime.datetime.now()
+        row = MyTable(created_on=created_on)
+        row.save().run_sync()
+
+        result = MyTable.objects().first().run_sync()
+        self.assertEqual(result.created_on, created_on)

tests/columns/test_uuid.py

@@ -16,7 +16,7 @@ def setUp(self):
     def tearDown(self):
         MyTable.alter().drop_table().run_sync()
 
-    def test_creation(self):
+    def test_return_type(self):
         row = MyTable()
         row.save().run_sync()