Update elasticsearch ext

Author: uriyyo

fastapi_pagination/bases.py

@@ -41,7 +41,7 @@ class PydanticUndefinedAnnotation(Exception):  # type: ignore[no-redef]
 
 from collections.abc import Sequence
 
-from typing_extensions import Self, TypeGuard
+from typing_extensions import Self, TypeIs
 
 from .types import Cursor, GreaterEqualZero, ParamsType
 
@@ -68,11 +68,11 @@ def as_cursor(self) -> CursorRawParams:
         raise ValueError("Not a 'cursor' params")
 
 
-def is_limit_offset(params: BaseRawParams) -> TypeGuard[RawParams]:
+def is_limit_offset(params: BaseRawParams) -> TypeIs[RawParams]:
     return params.type == "limit-offset"
 
 
-def is_cursor(params: BaseRawParams) -> TypeGuard[CursorRawParams]:
+def is_cursor(params: BaseRawParams) -> TypeIs[CursorRawParams]:
     return params.type == "cursor"
 
 

fastapi_pagination/ext/elasticsearch.py

@@ -2,54 +2,57 @@
     "paginate",
 ]
 
+from typing import Any, Optional
 
 from elasticsearch import Elasticsearch
 from elasticsearch_dsl import Search
 
 from fastapi_pagination.api import apply_items_transformer, create_page
-from fastapi_pagination.bases import AbstractParams, CursorRawParams, is_limit_offset
+from fastapi_pagination.bases import AbstractParams, is_limit_offset
 from fastapi_pagination.types import AdditionalData, SyncItemsTransformer
 from fastapi_pagination.utils import verify_params
 
 
 def paginate(
     conn: Elasticsearch,
     query: Search,
-    params: AbstractParams | None = None,
+    params: Optional[AbstractParams] = None,
     *,
-    transformer: SyncItemsTransformer | None = None,
-    additional_data: AdditionalData | None = None,
-):
+    transformer: Optional[SyncItemsTransformer] = None,
+    additional_data: Optional[AdditionalData] = None,
+) -> Any:
     params, raw_params = verify_params(params, "limit-offset", "cursor")
-    if is_limit_offset(raw_params):
-        total = query.using(conn).count()
-        response = query.using(conn)[raw_params.offset : raw_params.offset + raw_params.limit].execute()
-        items = response
-        t_items = apply_items_transformer(items, transformer)
-        return create_page(
-            t_items,
-            total=total,
-            params=params,
-            **(additional_data or {}),
-        )
 
     total = None
     if raw_params.include_total:
         total = query.using(conn).count()
-    raw_params: CursorRawParams
-    if not raw_params.cursor:
-        response = query.params(scroll="1m").extra(size=raw_params.size).execute()
-        items = response.hits
-        next_ = response._scroll_id
+
+    kwargs = {}
+    items: Any
+
+    if is_limit_offset(raw_params):
+        items = query.using(conn)[raw_params.as_slice()].execute()
     else:
-        response = conn.scroll(scroll_id=raw_params.cursor, scroll="1m")
-        next_ = response.get("_scroll_id")
-        items = [item.get("_source") for item in response["hits"]["hits"]]
+        raw_params = raw_params.as_cursor()
+
+        response: Any
+        if not raw_params.cursor:
+            response = query.params(scroll="1m").extra(size=raw_params.size).execute()
+            items = response.hits
+            next_ = response._scroll_id
+        else:
+            response = conn.scroll(scroll_id=raw_params.cursor, scroll="1m")  # type: ignore[arg-type]
+            next_ = response.get("_scroll_id")
+            items = [item.get("_source") for item in response["hits"]["hits"]]
+
+        kwargs["next_"] = next_
+
     t_items = apply_items_transformer(items, transformer)
+
     return create_page(
         t_items,
         total=total,
         params=params,
-        next_=next_,
+        **kwargs,
         **(additional_data or {}),
     )