diff --git a/README.md b/README.md index a34d88c..e0edf13 100644 --- a/README.md +++ b/README.md @@ -3,6 +3,7 @@ **Table of Contents** *generated with [DocToc](https://github.com/thlorenz/doctoc)* - [pynamodb-attributes](#pynamodb-attributes) + - [Testing](#testing) @@ -17,3 +18,16 @@ This Python 3 library contains compound and high-level PynamoDB attributes: - `TimestampAttribute`, `TimestampMsAttribute`, `TimestampUsAttribute` – serializes `datetime`s as Unix epoch seconds, milliseconds (ms) or microseconds (µs) - `IntegerDateAttribute` - serializes `date` as an integer representing the Gregorian date (_e.g._ `20181231`) - `UUIDAttribute` - serializes a `UUID` Python object as a `S` type attribute (_e.g._ `'a8098c1a-f86e-11da-bd1a-00112444be1e'`) +- `UnicodeDatetimeAttribute` - ISO8601 datetime strings with offset information + +## Testing + +The tests in this repository use an in-memory implementation of [`dynamodb`](https://aws.amazon.com/dynamodb). To run the tests locally, make sure [DynamoDB Local](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.html) is running. It is available as a standalone binary, through package managers (e.g. [Homebrew](https://formulae.brew.sh/cask/dynamodb-local)) or as a Docker container: +```shell +docker run -d -p 8000:8000 amazon/dynamodb-local +``` + +Afterwards, run tests as usual: +```shell +pytest tests +``` diff --git a/pynamodb_attributes/__init__.py b/pynamodb_attributes/__init__.py index 8d72c72..69522c5 100644 --- a/pynamodb_attributes/__init__.py +++ b/pynamodb_attributes/__init__.py @@ -8,6 +8,7 @@ from .timestamp import TimestampAttribute from .timestamp import TimestampMsAttribute from .timestamp import TimestampUsAttribute +from .unicode_datetime import UnicodeDatetimeAttribute from .unicode_delimited_tuple import UnicodeDelimitedTupleAttribute from .unicode_enum import UnicodeEnumAttribute from .uuid import UUIDAttribute @@ -26,4 +27,5 @@ "TimestampMsAttribute", "TimestampUsAttribute", "UUIDAttribute", + "UnicodeDatetimeAttribute", ] diff --git a/pynamodb_attributes/unicode_datetime.py b/pynamodb_attributes/unicode_datetime.py new file mode 100644 index 0000000..8c46255 --- /dev/null +++ b/pynamodb_attributes/unicode_datetime.py @@ -0,0 +1,55 @@ +from datetime import datetime +from datetime import timezone +from typing import Any +from typing import Optional + +import pynamodb.attributes +from pynamodb.attributes import Attribute + + +class UnicodeDatetimeAttribute(Attribute[datetime]): + """ + Stores a 'datetime.datetime' object as an ISO8601 formatted string + + This is useful for wanting database readable datetime objects that also sort. + + >>> class MyModel(Model): + >>> created_at = UnicodeDatetimeAttribute() + """ + + attr_type = pynamodb.attributes.STRING + + def __init__( + self, + *, + force_tz: bool = True, + force_utc: bool = False, + fmt: Optional[str] = None, + **kwargs: Any, + ) -> None: + """ + :param force_tz: If set it will add timezone info to the `datetime` value if no `tzinfo` is currently + set before serializing, defaults to `True` + :param force_utc: If set it will normalize the `datetime` to UTC before serializing the value + :param fmt: If set this value will be used to format the `datetime` object for serialization + and deserialization + """ + + super().__init__(**kwargs) + self._force_tz = force_tz + self._force_utc = force_utc + self._fmt = fmt + + def deserialize(self, value: str) -> datetime: + return ( + datetime.fromisoformat(value) + if self._fmt is None + else datetime.strptime(value, self._fmt) + ) + + def serialize(self, value: datetime) -> str: + if self._force_tz and value.tzinfo is None: + value = value.replace(tzinfo=timezone.utc) + if self._force_utc: + value = value.astimezone(tz=timezone.utc) + return value.isoformat() if self._fmt is None else value.strftime(self._fmt) diff --git a/setup.py b/setup.py index 83ab80f..c49ff3e 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ setup( name="pynamodb-attributes", - version="0.3.0", + version="0.3.1", description="Common attributes for PynamoDB", url="https://www.github.com/lyft/pynamodb-attributes", maintainer="Lyft", diff --git a/tests/mypy_test.py b/tests/mypy_test.py index 12b7858..5720973 100644 --- a/tests/mypy_test.py +++ b/tests/mypy_test.py @@ -141,3 +141,18 @@ class MyModel(Model): reveal_type(MyModel().my_attr) # E: Revealed type is "uuid.UUID*" """, ) + + +def test_unicode_datetime(): + assert_mypy_output( + """ + from pynamodb.models import Model + from pynamodb_attributes import UnicodeDatetimeAttribute + + class MyModel(Model): + my_attr = UnicodeDatetimeAttribute() + + reveal_type(MyModel.my_attr) # E: Revealed type is "pynamodb_attributes.unicode_datetime.UnicodeDatetimeAttribute" + reveal_type(MyModel().my_attr) # E: Revealed type is "datetime.datetime*" + """, + ) diff --git a/tests/unicode_datetime_attribute_test.py b/tests/unicode_datetime_attribute_test.py new file mode 100644 index 0000000..ffe9032 --- /dev/null +++ b/tests/unicode_datetime_attribute_test.py @@ -0,0 +1,206 @@ +from datetime import datetime +from unittest.mock import ANY + +import pytest +from pynamodb.attributes import UnicodeAttribute +from pynamodb.models import Model + +from pynamodb_attributes import UnicodeDatetimeAttribute +from tests.connection import _connection +from tests.meta import dynamodb_table_meta + + +CUSTOM_FORMAT = "%m/%d/%Y, %H:%M:%S" +CUSTOM_FORMAT_DATE = "11/22/2020, 11:22:33" +TEST_ISO_DATE_NO_OFFSET = "2020-11-22T11:22:33.444444" +TEST_ISO_DATE_UTC = "2020-11-22T11:22:33.444444+00:00" +TEST_ISO_DATE_PST = "2020-11-22T03:22:33.444444-08:00" + + +class MyModel(Model): + Meta = dynamodb_table_meta(__name__) + + key = UnicodeAttribute(hash_key=True) + default = UnicodeDatetimeAttribute(null=True) + no_force_tz = UnicodeDatetimeAttribute(force_tz=False, null=True) + force_utc = UnicodeDatetimeAttribute(force_utc=True, null=True) + force_utc_no_force_tz = UnicodeDatetimeAttribute( + force_utc=True, + force_tz=False, + null=True, + ) + custom_format = UnicodeDatetimeAttribute(fmt=CUSTOM_FORMAT, null=True) + + +@pytest.fixture(scope="module", autouse=True) +def create_table(): + MyModel.create_table() + + +@pytest.mark.parametrize( + ["value", "expected_str", "expected_value"], + [ + ( + datetime.fromisoformat(TEST_ISO_DATE_NO_OFFSET), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_UTC), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_PST), + TEST_ISO_DATE_PST, + datetime.fromisoformat(TEST_ISO_DATE_PST), + ), + ], +) +def test_default_serialization(value, expected_str, expected_value, uuid_key): + model = MyModel() + model.key = uuid_key + model.default = value + + model.save() + + actual = MyModel.get(hash_key=uuid_key) + assert actual.default == expected_value + + item = _connection(MyModel).get_item(uuid_key) + assert item["Item"] == {"key": ANY, "default": {"S": expected_str}} + + +@pytest.mark.parametrize( + ["value", "expected_str", "expected_value"], + [ + ( + datetime.fromisoformat(TEST_ISO_DATE_NO_OFFSET), + TEST_ISO_DATE_NO_OFFSET, + datetime.fromisoformat(TEST_ISO_DATE_NO_OFFSET), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_UTC), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_PST), + TEST_ISO_DATE_PST, + datetime.fromisoformat(TEST_ISO_DATE_PST), + ), + ], +) +def test_no_force_tz_serialization(value, expected_str, expected_value, uuid_key): + model = MyModel() + model.key = uuid_key + model.no_force_tz = value + + model.save() + + actual = MyModel.get(hash_key=uuid_key) + item = _connection(MyModel).get_item(uuid_key) + + assert item["Item"] == {"key": ANY, "no_force_tz": {"S": expected_str}} + + assert actual.no_force_tz == expected_value + + +@pytest.mark.parametrize( + ["value", "expected_str", "expected_value"], + [ + ( + datetime.fromisoformat(TEST_ISO_DATE_NO_OFFSET), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_UTC), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_PST), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ], +) +def test_force_utc_serialization(value, expected_str, expected_value, uuid_key): + model = MyModel() + model.key = uuid_key + model.force_utc = value + + model.save() + + actual = MyModel.get(hash_key=uuid_key) + item = _connection(MyModel).get_item(uuid_key) + + assert item["Item"] == {"key": ANY, "force_utc": {"S": expected_str}} + + assert actual.force_utc == expected_value + + +@pytest.mark.parametrize( + ["value", "expected_str", "expected_value"], + [ + ( + datetime.fromisoformat(TEST_ISO_DATE_UTC), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ( + datetime.fromisoformat(TEST_ISO_DATE_PST), + TEST_ISO_DATE_UTC, + datetime.fromisoformat(TEST_ISO_DATE_UTC), + ), + ], +) +def test_force_utc_no_force_tz_serialization( + value, + expected_str, + expected_value, + uuid_key, +): + model = MyModel() + model.key = uuid_key + model.force_utc_no_force_tz = value + + model.save() + + actual = MyModel.get(hash_key=uuid_key) + item = _connection(MyModel).get_item(uuid_key) + + assert item["Item"] == {"key": ANY, "force_utc_no_force_tz": {"S": expected_str}} + + assert actual.force_utc_no_force_tz == expected_value + + +@pytest.mark.parametrize( + ["value", "expected_str", "expected_value"], + [ + ( + datetime.fromisoformat(TEST_ISO_DATE_UTC), + CUSTOM_FORMAT_DATE, + datetime(2020, 11, 22, 11, 22, 33), + ), + ], +) +def test_custom_format_force_tz_serialization( + value, + expected_str, + expected_value, + uuid_key, +): + model = MyModel() + model.key = uuid_key + model.custom_format = value + + model.save() + + actual = MyModel.get(hash_key=uuid_key) + item = _connection(MyModel).get_item(uuid_key) + + assert item["Item"] == {"key": ANY, "custom_format": {"S": expected_str}} + + assert actual.custom_format == expected_value