feat: adds uuid_utils type decoder automatically (#145)

litestar-org · Mar 20, 2024 · e3b4652 · e3b4652
1 parent 2a8bf18
commit e3b4652
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/advanced_alchemy/base.py b/advanced_alchemy/base.py
@@ -21,7 +21,9 @@
 from advanced_alchemy.types import GUID, UUID_UTILS_INSTALLED, BigIntIdentity, DateTimeUTC, JsonB
 
 if UUID_UTILS_INSTALLED and not TYPE_CHECKING:
-    from uuid_utils import UUID, uuid4, uuid6, uuid7  # pyright: ignore[reportMissingImports]
+    from uuid_utils import UUID  # pyright: ignore[reportMissingImports]
+    from uuid_utils.compat import uuid4, uuid6, uuid7  # pyright: ignore[reportMissingImports]
+
 else:
     from uuid import UUID, uuid4  # type: ignore[assignment]
 

diff --git a/advanced_alchemy/extensions/litestar/plugins/init/plugin.py b/advanced_alchemy/extensions/litestar/plugins/init/plugin.py
@@ -4,6 +4,9 @@
 from typing import TYPE_CHECKING
 
 from litestar.di import Provide
+from litestar.dto import DTOData
+from litestar.pagination import ClassicPagination, CursorPagination, OffsetPagination
+from litestar.params import Dependency, Parameter
 from litestar.plugins import CLIPluginProtocol, InitPluginProtocol
 
 from advanced_alchemy.extensions.litestar.plugins import _slots_base
@@ -37,6 +40,12 @@
     "NotInCollectionFilter": NotInCollectionFilter,
     "NotInSearchFilter": NotInSearchFilter,
     "FilterTypes": FilterTypes,
+    "OffsetPagination": OffsetPagination,
+    "ClassicPagination": ClassicPagination,
+    "CursorPagination": CursorPagination,
+    "Parameter": Parameter,
+    "Dependency": Dependency,
+    "DTOData": DTOData,
 }
 
 
@@ -70,13 +79,14 @@ def on_app_init(self, app_config: AppConfig) -> AppConfig:
             signature_namespace_values.update({"pgproto.UUID": pgproto.UUID})
             app_config.type_encoders = {pgproto.UUID: str, **(app_config.type_encoders or {})}
         with contextlib.suppress(ImportError):
-            from uuid_utils import UUID  # pyright: ignore[reportMissingImports]
-
-            signature_namespace_values.update({"UUID": UUID})
-            app_config.type_encoders = {
-                **(app_config.type_encoders or {}),
-                UUID: str,
-            }
+            import uuid_utils  # pyright: ignore[reportMissingImports]
+
+            signature_namespace_values.update({"uuid_utils.UUID": uuid_utils.UUID})
+            app_config.type_encoders = {uuid_utils.UUID: str, **(app_config.type_encoders or {})}
+            app_config.type_decoders = [
+                (lambda x: x is uuid_utils.UUID, lambda t, v: t(str(v))),
+                *(app_config.type_decoders or []),
+            ]
         app_config.signature_namespace.update(self._config.signature_namespace)
         app_config.signature_namespace.update(signature_namespace_values)
         app_config.lifespan.append(self._config.lifespan)