Support time points as numeric indexes

aredis_om/model/encoders.py

@@ -23,8 +23,11 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
-
+import calendar
+import copy
 import dataclasses
+import datetime
+import math
 from collections import defaultdict
 from enum import Enum
 from pathlib import PurePath
@@ -35,8 +38,36 @@
 from pydantic.json import ENCODERS_BY_TYPE
 
 
+# TODO: check if correct
+def date_to_timestamp(t: datetime.date) -> int:
+    return calendar.timegm(t.timetuple())
+
+
+# TODO: check if correct
+def datetime_to_timestamp(t: datetime.datetime) -> int:
+    return math.floor(t.astimezone(datetime.timezone.utc).timestamp() * 1000)
+
+
+zero_time = datetime.datetime.fromtimestamp(0)
+zero_day = zero_time.date()
+
+
+# TODO: check if correct
+def time_to_timestamp(t: datetime.time) -> int:
+    time_point = datetime.datetime.combine(zero_day, t, t.tzinfo)
+    return datetime_to_timestamp(time_point)
+
+
 SetIntStr = Set[Union[int, str]]
 DictIntStrAny = Dict[Union[int, str], Any]
+ENCODERS_BY_TYPE_ENHANCED = copy.copy(ENCODERS_BY_TYPE)
+ENCODERS_BY_TYPE_ENHANCED.update(
+    {
+        datetime.date: date_to_timestamp,
+        datetime.datetime: datetime_to_timestamp,
+        datetime.time: time_to_timestamp,
+    }
+)
 
 
 def generate_encoders_by_class_tuples(
@@ -154,8 +185,8 @@ def jsonable_encoder(
                 if isinstance(obj, encoder_type):
                     return encoder(obj)
 
-    if type(obj) in ENCODERS_BY_TYPE:
-        return ENCODERS_BY_TYPE[type(obj)](obj)
+    if type(obj) in ENCODERS_BY_TYPE_ENHANCED:
+        return ENCODERS_BY_TYPE_ENHANCED[type(obj)](obj)
     for encoder, classes_tuple in encoders_by_class_tuples.items():
         if isinstance(obj, classes_tuple):
             return encoder(obj)

aredis_om/model/model.py

@@ -1,5 +1,6 @@
 import abc
 import dataclasses
+import datetime
 import decimal
 import json
 import logging
@@ -38,7 +39,12 @@
 from ..checks import has_redis_json, has_redisearch
 from ..connections import get_redis_connection
 from ..unasync_util import ASYNC_MODE
-from .encoders import jsonable_encoder
+from .encoders import (
+    date_to_timestamp,
+    datetime_to_timestamp,
+    jsonable_encoder,
+    time_to_timestamp,
+)
 from .render_tree import render_tree
 from .token_escaper import TokenEscaper
 
@@ -330,7 +336,14 @@ class RediSearchFieldTypes(Enum):
 
 
 # TODO: How to handle Geo fields?
-NUMERIC_TYPES = (float, int, decimal.Decimal)
+NUMERIC_TYPES = (
+    float,
+    int,
+    decimal.Decimal,
+    datetime.date,
+    datetime.datetime,
+    datetime.time,
+)
 DEFAULT_PAGE_SIZE = 1000
 
 
@@ -530,6 +543,7 @@ def resolve_value(
                     f"Docs: {ERRORS_URL}#E5"
                 )
         elif field_type is RediSearchFieldTypes.NUMERIC:
+            value = jsonable_encoder(value)
             if op is Operators.EQ:
                 result += f"@{field_name}:[{value} {value}]"
             elif op is Operators.NE:
@@ -1461,6 +1475,13 @@ def schema_for_type(cls, name, typ: Any, field_info: PydanticFieldInfo):
 
 
 class JsonModel(RedisModel, abc.ABC):
+    class Config(RedisModel.Config):
+        json_encoders = {
+            datetime.date: date_to_timestamp,
+            datetime.datetime: datetime_to_timestamp,
+            datetime.time: time_to_timestamp,
+        }
+
     def __init_subclass__(cls, **kwargs):
         # Generate the RediSearch schema once to validate fields.
         cls.redisearch_schema()

tests/conftest.py

@@ -1,5 +1,6 @@
 import asyncio
 import random
+from typing import Iterable
 
 import pytest
 
@@ -37,6 +38,20 @@ def _delete_test_keys(prefix: str, conn):
         conn.delete(*keys)
 
 
+def _delete_test_indexes(prefix: str, conn):
+    # TODO: move to scan when available
+    # https://redis.io/commands/ft._list/
+    from redis import ResponseError
+
+    try:
+        indexes: Iterable[str] = conn.execute_command("ft._list")
+    except ResponseError:
+        return
+    for index in indexes:
+        if index.startswith(prefix):
+            conn.execute_command("ft.dropindex", index, "dd")
+
+
 @pytest.fixture
 def key_prefix(request, redis):
     key_prefix = f"{TEST_PREFIX}:{random.random()}"
@@ -59,3 +74,4 @@ def cleanup_keys(request):
     # Delete keys only once
     if conn.decr(once_key) == 0:
         _delete_test_keys(TEST_PREFIX, conn)
+        _delete_test_indexes(TEST_PREFIX, conn)

tests/test_hash_model.py

@@ -25,6 +25,7 @@
 from redis_om import has_redisearch
 from tests.conftest import py_test_mark_asyncio
 
+
 if not has_redisearch():
     pytestmark = pytest.mark.skip
 

tests/test_json_model.py

@@ -27,6 +27,7 @@
 from redis_om import has_redis_json
 from tests.conftest import py_test_mark_asyncio
 
+
 if not has_redis_json():
     pytestmark = pytest.mark.skip
 
@@ -437,7 +438,11 @@ async def test_recursive_query_expression_resolution(members, m):
 async def test_recursive_query_field_resolution(members, m):
     member1, _, _ = members
     member1.address.note = m.Note(
-        description="Weird house", created_on=datetime.datetime.now()
+        description="Weird house",
+        created_on=datetime.datetime.now().replace(
+            microsecond=0,
+            tzinfo=datetime.timezone.utc,
+        ),
     )
     await member1.save()
     actual = await m.Member.find(
@@ -449,7 +454,10 @@ async def test_recursive_query_field_resolution(members, m):
         m.Order(
             items=[m.Item(price=10.99, name="Ball")],
             total=10.99,
-            created_on=datetime.datetime.now(),
+            created_on=datetime.datetime.now().replace(
+                microsecond=0,
+                tzinfo=datetime.timezone.utc,
+            ),
         )
     ]
     await member1.save()