Add pydantic v2

sopherapps · Feb 10, 2024 · bbf8b1e · bbf8b1e
1 parent 298283e
commit bbf8b1e
Showing 11 changed files with 111 additions and 71 deletions.
diff --git a/README.md b/README.md
@@ -24,6 +24,7 @@ Most Notable Features are:
 
 ## Benchmarks
 
+### v0.
 On an average PC ~16GB RAM, i7 Core
 
 ```
@@ -55,6 +56,37 @@ benchmark_bulk_insert[redis_store]                         1,025.0436 (8.29)
 -------------------------------------------------------------------------------------------------------------------------
 ```
 
+# v1. (with pydantic v2)
+
+```
+------------------------------------------------- benchmark: 22 tests -------------------------------------------------
+Name (time in us)                                              Mean                 Min                   Max          
+-----------------------------------------------------------------------------------------------------------------------
+benchmark_delete[redis_store-Wuthering Heights]            124.1668 (1.0)      108.9610 (1.0)        418.1310 (1.04)   
+benchmark_bulk_delete[redis_store]                         137.7564 (1.11)     121.6380 (1.12)       470.7510 (1.17)   
+benchmark_select_columns_for_one_id[redis_store-book2]     166.7328 (1.34)     147.9490 (1.36)       430.3780 (1.07)   
+benchmark_select_columns_for_one_id[redis_store-book1]     171.0826 (1.38)     148.6430 (1.36)       426.0820 (1.06)   
+benchmark_select_columns_for_one_id[redis_store-book0]     171.7202 (1.38)     148.6460 (1.36)       431.3730 (1.07)   
+benchmark_select_columns_for_one_id[redis_store-book3]     172.1800 (1.39)     148.9410 (1.37)       471.5910 (1.17)   
+benchmark_select_all_for_one_id[redis_store-book1]         189.0068 (1.52)     163.5860 (1.50)       457.3090 (1.14)   
+benchmark_select_all_for_one_id[redis_store-book2]         188.5258 (1.52)     163.6650 (1.50)       401.7030 (1.0)    
+benchmark_select_all_for_one_id[redis_store-book3]         187.5434 (1.51)     165.3890 (1.52)       460.7100 (1.15)   
+benchmark_select_all_for_one_id[redis_store-book0]         190.3049 (1.53)     165.7280 (1.52)       459.8080 (1.14)   
+benchmark_select_columns_for_some_items[redis_store]       222.1405 (1.79)     198.9940 (1.83)       485.6230 (1.21)   
+benchmark_select_columns_paginated[redis_store]            229.5429 (1.85)     200.4560 (1.84)       494.4250 (1.23)   
+benchmark_select_default_paginated[redis_store]            262.3155 (2.11)     231.3960 (2.12)       568.8410 (1.42)   
+benchmark_select_some_items[redis_store]                   270.4251 (2.18)     232.3230 (2.13)       537.2130 (1.34)   
+benchmark_update[redis_store-Wuthering Heights-data0]      280.6308 (2.26)     248.7310 (2.28)       676.0330 (1.68)   
+benchmark_select_columns[redis_store]                      316.7642 (2.55)     283.6720 (2.60)       560.9610 (1.40)   
+benchmark_single_insert[redis_store-book2]                 343.9583 (2.77)     284.1000 (2.61)       585.6200 (1.46)   
+benchmark_single_insert[redis_store-book1]                 328.5308 (2.65)     291.8760 (2.68)       600.8130 (1.50)   
+benchmark_single_insert[redis_store-book3]                 341.0249 (2.75)     292.2800 (2.68)       575.1020 (1.43)   
+benchmark_single_insert[redis_store-book0]                 349.9540 (2.82)     299.6660 (2.75)       606.0370 (1.51)   
+benchmark_select_default[redis_store]                      381.0231 (3.07)     346.2910 (3.18)       669.7460 (1.67)   
+benchmark_bulk_insert[redis_store]                         840.4876 (6.77)     790.2340 (7.25)     1,049.8260 (2.61)   
+-----------------------------------------------------------------------------------------------------------------------
+```
+
 ## Contributions
 
 Contributions are welcome. The docs have to maintained, the code has to be made cleaner, more idiomatic and faster,

diff --git a/pydantic_redis/_shared/model/base.py b/pydantic_redis/_shared/model/base.py
@@ -4,7 +4,8 @@
 import typing
 from typing import Dict, Tuple, Any, Type, Union, List, Optional
 
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
+from pydantic.fields import ModelPrivateAttr
 
 from pydantic_redis._shared.utils import (
     typing_get_origin,
@@ -45,9 +46,7 @@ class AbstractModel(BaseModel):
     _nested_model_tuple_fields: Dict[str, Tuple[Any, ...]] = {}
     _nested_model_list_fields: Dict[str, Type["AbstractModel"]] = {}
     _nested_model_fields: Dict[str, Type["AbstractModel"]] = {}
-
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @classmethod
     def get_store(cls) -> AbstractStore:
@@ -93,7 +92,10 @@ def get_primary_key_field(cls):
         Returns:
             the field that can be used to uniquely identify each record of current Model
         """
-        return cls._primary_key_field
+        try:
+            return cls._primary_key_field.get_default()
+        except AttributeError:
+            return cls._primary_key_field
 
     @classmethod
     def get_field_types(cls) -> Dict[str, Any]:

diff --git a/pydantic_redis/_shared/store.py b/pydantic_redis/_shared/store.py
@@ -3,9 +3,10 @@
 """
 from typing import Optional, Union, Type, Dict, Any
 
+from pydantic.fields import ModelPrivateAttr
 from redis import Redis
 from redis.asyncio import Redis as AioRedis
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 from redis.commands.core import Script, AsyncScript
 
 from ..config import RedisConfig
@@ -45,10 +46,7 @@ class AbstractStore(BaseModel):
     ] = None
     select_some_fields_for_some_ids_script: Optional[Union[AsyncScript, Script]] = None
     models: Dict[str, Type["AbstractModel"]] = {}
-
-    class Config:
-        arbitrary_types_allowed = True
-        orm_mode = True
+    model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)
 
     def __init__(
         self,
@@ -122,7 +120,7 @@ def register_model(self, model_class: Type["AbstractModel"]):
                 a certain type of records to be saved in redis.
         """
         if not isinstance(model_class.get_primary_key_field(), str):
-            raise NotImplementedError(
+            raise AttributeError(
                 f"{model_class.__name__} should have a _primary_key_field"
             )
 

diff --git a/pydantic_redis/_shared/utils.py b/pydantic_redis/_shared/utils.py
@@ -2,7 +2,7 @@
 
 """
 import typing
-from typing import Any, Tuple, Optional, Union, Dict, Callable, Type, List
+from typing import Any, Tuple, Optional, Union, Dict, Type, List
 
 import orjson
 
@@ -63,15 +63,16 @@ def from_bytes_to_str(value: Union[str, bytes]) -> str:
     Returns:
         the string value of the argument passed
     """
-    if isinstance(value, bytes):
+    try:
         return str(value, "utf-8")
-    return value
+    except TypeError:
+        return value
 
 
 def from_str_or_bytes_to_any(value: Any, field_type: Type) -> Any:
     """Converts str or bytes to arbitrary data.
 
-    Converts the the `value` from a string or bytes to the `field_type`.
+    Converts the `value` from a string or bytes to the `field_type`.
 
     Args:
         value: the string or bytes to be transformed to the `field_type`
@@ -116,9 +117,10 @@ def default_json_dump(obj: Any):
     Returns:
         the bytes or string value of the object
     """
-    if hasattr(obj, "json") and isinstance(obj.json, Callable):
-        return obj.json()
-    return obj
+    try:
+        return obj.model_dump_json()
+    except AttributeError:
+        return obj
 
 
 def from_dict_to_key_value_list(data: Dict[str, Any]) -> List[Any]:

diff --git a/pydantic_redis/asyncio/model.py b/pydantic_redis/asyncio/model.py
@@ -184,3 +184,6 @@ async def select(
         return parse_select_response(
             model=cls, response=response, as_models=(columns is None)
         )
+
+
+Store.model_rebuild()
diff --git a/pydantic_redis/config.py b/pydantic_redis/config.py
@@ -2,7 +2,7 @@
 """
 from typing import Optional
 
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 
 
 class RedisConfig(BaseModel):
@@ -24,6 +24,8 @@ class RedisConfig(BaseModel):
             (default: utf-8)
     """
 
+    model_config = ConfigDict(from_attributes=True)
+
     host: str = "localhost"
     port: int = 6379
     db: int = 0
@@ -38,6 +40,3 @@ def redis_url(self) -> str:
         if self.password is None:
             return f"{proto}://{self.host}:{self.port}/{self.db}"
         return f"{proto}://:{self.password}@{self.host}:{self.port}/{self.db}"
-
-    class Config:
-        orm_mode = True
diff --git a/pydantic_redis/syncio/model.py b/pydantic_redis/syncio/model.py
@@ -180,3 +180,6 @@ def select(
         return parse_select_response(
             model=cls, response=response, as_models=(columns is None)
         )
+
+
+Store.model_rebuild()
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 orjson==3.6.1
 hiredis==2.0.0
-pydantic==1.9.2
+pydantic==2.6.1
 pytest==7.0.1
 pytest-benchmark==3.4.1
 pytest-lazy-fixture==0.6.3

diff --git a/test/conftest.py b/test/conftest.py
@@ -50,21 +50,21 @@ class Library(syn.Model):
     _primary_key_field: str = "name"
     name: str
     address: str
-    books: List[Book] = None
+    books: List[Book] = []
     lost: Optional[List[Book]] = None
     popular: Optional[Tuple[Book, Book]] = None
-    new: Tuple[Book, Author, Book, int] = None
+    new: Optional[Tuple[Book, Author, Book, int]] = None
 
 
 class AsyncLibrary(asy.Model):
     # the _primary_key_field is mandatory
     _primary_key_field: str = "name"
     name: str
     address: str
-    books: List[AsyncBook] = None
+    books: List[AsyncBook] = []
     lost: Optional[List[AsyncBook]] = None
     popular: Optional[Tuple[AsyncBook, AsyncBook]] = None
-    new: Tuple[AsyncBook, AsyncAuthor, AsyncBook, int] = None
+    new: Optional[Tuple[AsyncBook, AsyncAuthor, AsyncBook, int]] = None
 
 
 authors = {
@@ -113,30 +113,30 @@ class AsyncLibrary(asy.Model):
 async_books = [
     AsyncBook(
         title="Oliver Twist",
-        author=authors["charles"],
+        author=async_authors["charles"],
         published_on=date(year=1215, month=4, day=4),
         in_stock=False,
         rating=2,
         tags=["Classic"],
     ),
     AsyncBook(
         title="Great Expectations",
-        author=authors["charles"],
+        author=async_authors["charles"],
         published_on=date(year=1220, month=4, day=4),
         rating=5,
         tags=["Classic"],
     ),
     AsyncBook(
         title="Jane Eyre",
-        author=authors["charles"],
+        author=async_authors["charles"],
         published_on=date(year=1225, month=6, day=4),
         in_stock=False,
         rating=3.4,
         tags=["Classic", "Romance"],
     ),
     AsyncBook(
         title="Wuthering Heights",
-        author=authors["jane"],
+        author=async_authors["jane"],
         published_on=date(year=1600, month=4, day=4),
         rating=4.0,
         tags=["Classic", "Romance"],