refactor: add generic type parameters for improved type flexibility

src/typesense/__init__.py

@@ -1,4 +1,4 @@
 from .client import Client  # NOQA
 
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"

src/typesense/types/collection.py

@@ -10,6 +10,22 @@
 
 _TType = typing.TypeVar("_TType")
 
+_CreateFieldT = typing.TypeVar(
+    "_CreateFieldT",
+    bound=typing.Union[
+        "RegularCollectionFieldSchema", "ReferenceCollectionFieldSchema"
+    ],
+)
+
+_UpdateFieldT = typing.TypeVar(
+    "_UpdateFieldT",
+    bound=typing.Union[
+        "RegularCollectionFieldSchema",
+        "ReferenceCollectionFieldSchema",
+        "DropCollectionFieldSchema",
+    ],
+)
+
 _FieldType = typing.Literal[
     "string",
     "int32",
@@ -150,15 +166,14 @@ class VoiceQueryModelSchema(typing.TypedDict):
     model_name: str
 
 
-class CollectionCreateSchema(typing.TypedDict):
+class CollectionCreateSchema(typing.Generic[_CreateFieldT], typing.TypedDict):
     """
     The schema for the request of the Collections.create method.
 
     Attributes:
         name (str): The name of the collection.
 
-        fields (list[RegularCollectionFieldSchema | ReferenceCollectionFieldSchema]): The fields
-            of the collection.
+        fields (list[_CreateFieldT]): The fields of the collection.
 
         default_sorting_field (str): The default sorting field of the collection.
 
@@ -172,17 +187,15 @@ class CollectionCreateSchema(typing.TypedDict):
     """
 
     name: str
-    fields: typing.List[
-        typing.Union[RegularCollectionFieldSchema, ReferenceCollectionFieldSchema]
-    ]
+    fields: typing.List[_CreateFieldT]
     default_sorting_field: typing.NotRequired[str]
     symbols_to_index: typing.NotRequired[typing.List[str]]
     token_separators: typing.NotRequired[typing.List[str]]
     enable_nested_fields: typing.NotRequired[bool]
     voice_query_model: typing.NotRequired[VoiceQueryModelSchema]
 
 
-class CollectionSchema(CollectionCreateSchema):
+class CollectionSchema(CollectionCreateSchema[_CreateFieldT]):
     """
     The schema for the response of the Collections.create method.
 
@@ -195,8 +208,7 @@ class CollectionSchema(CollectionCreateSchema):
 
         name (str): The name of the collection.
 
-        fields (list[RegularCollectionFieldSchema | ReferenceCollectionFieldSchema]): The fields
-            of the collection.
+        fields (list[_CreateFieldT]): The fields of the collection.
 
         default_sorting_field (str): The default sorting field of the collection.
 
@@ -214,19 +226,29 @@ class CollectionSchema(CollectionCreateSchema):
     num_memory_shards: int
 
 
-class CollectionUpdateSchema(typing.TypedDict):
+# Type alias for backward compatibility
+CollectionSchemaCompat = CollectionSchema[
+    typing.Union[RegularCollectionFieldSchema, ReferenceCollectionFieldSchema]
+]
+
+
+class CollectionUpdateSchema(typing.Generic[_UpdateFieldT], typing.TypedDict):
     """
     The schema for the request of the Collection.update method.
 
     Attributes:
-        fields (list): The fields of the collection.
+        fields (list[_UpdateFieldT]): The fields of the collection.
 
     """
 
-    fields: typing.List[
-        typing.Union[
-            RegularCollectionFieldSchema,
-            ReferenceCollectionFieldSchema,
-            DropCollectionFieldSchema,
-        ]
+    fields: typing.List[_UpdateFieldT]
+
+
+# Type alias for backward compatibility
+CollectionUpdateSchemaCompat = CollectionUpdateSchema[
+    typing.Union[
+        RegularCollectionFieldSchema,
+        ReferenceCollectionFieldSchema,
+        DropCollectionFieldSchema,
     ]
+]

src/typesense/types/document.py

@@ -187,6 +187,33 @@ class ImportResponseFail(typing.Generic[TDoc], typing.TypedDict):
     document: TDoc
 
 
+_DocumentImportParamsT = typing.TypeVar(
+    "_DocumentImportParamsT",
+    bound=typing.Union[
+        "DocumentWriteParameters",
+        "DocumentImportParametersReturnId",
+        "DocumentImportParametersReturnDoc",
+        "DocumentImportParametersReturnDocAndId",
+    ],
+)
+
+_ImportResponseT = typing.TypeVar(
+    "_ImportResponseT",
+    bound=typing.Union[
+        "ImportResponseSuccess",
+        "ImportResponseWithDoc[typing.Any]",
+        "ImportResponseWithId",
+        "ImportResponseWithDocAndId[typing.Any]",
+        "ImportResponseFail[typing.Any]",
+    ],
+)
+
+_StringOrListT = typing.TypeVar(
+    "_StringOrListT",
+    bound=typing.Union[str, typing.List[str]],
+)
+
+
 ImportResponse: typing.TypeAlias = typing.Union[
     typing.List[typing.Union[ImportResponseWithDoc[TDoc], ImportResponseFail[TDoc]]],
     typing.List[typing.Union[ImportResponseWithId, ImportResponseFail[TDoc]]],
@@ -299,6 +326,16 @@ class DocumentImportParametersReturnDocAndId(DocumentWriteParameters):
     return_id: typing.Literal[True]
 
 
+DocumentImportParametersGeneric = typing.TypeVar(
+    "DocumentImportParametersGeneric",
+    bound=typing.Union[
+        DocumentWriteParameters,
+        DocumentImportParametersReturnId,
+        DocumentImportParametersReturnDoc,
+        DocumentImportParametersReturnDocAndId,
+    ],
+)
+
 DocumentImportParameters: typing.TypeAlias = typing.Union[
     DocumentWriteParameters,
     DocumentImportParametersReturnId,

src/typesense/types/key.py

@@ -8,6 +8,20 @@
     import typing_extensions as typing
 
 
+_ActionT = typing.TypeVar(
+    "_ActionT",
+    bound=typing.Union[
+        "_CollectionActions",
+        "_DocumentActions",
+        "_AliasActions",
+        "_SynonymActions",
+        "_OverrideActions",
+        "_StopwordActions",
+        "_KeyActions",
+        "_MiscActions",
+    ],
+)
+
 _CollectionActions = typing.Literal[
     "collections:list",
     "collections:get",
@@ -87,28 +101,31 @@
 ]
 
 
-class ApiKeyCreateSchema(typing.TypedDict):
+class ApiKeyCreateSchema(typing.Generic[_ActionT], typing.TypedDict):
     """
     Schema for creating a [new API key](https://typesense.org/docs/26.0/api/api-keys.html#create-an-api-key).
 
     Attributes:
-        actions (list[str]): The actions allowed for this key.
+        actions (list[_ActionT]): The actions allowed for this key.
         collections (list[str]): The collections this key has access to.
         description (str): The description for this key.
         value (str): The value of the key.
         expires_at (int): The time in UNIX timestamp format when the key will expire.
         autodelete (bool): Whether the key should be deleted after it expires.
     """
 
-    actions: typing.List[_Actions]
+    actions: typing.List[_ActionT]
     collections: typing.List[str]
     description: str
     value: typing.NotRequired[str]
     expires_at: typing.NotRequired[int]
     autodelete: typing.NotRequired[bool]
 
 
-class ApiKeyCreateResponseSchema(ApiKeyCreateSchema):
+ApiKeyCreateSchemaCompat = ApiKeyCreateSchema[_Actions]
+
+
+class ApiKeyCreateResponseSchema(ApiKeyCreateSchema[_ActionT]):
     """
     Response schema for creating a [new API key](https://typesense.org/docs/26.0/api/api-keys.html#create-an-api-key).
 
@@ -121,36 +138,45 @@ class ApiKeyCreateResponseSchema(ApiKeyCreateSchema):
     id: int
 
 
-class ApiKeySchema(typing.TypedDict):
+ApiKeyCreateResponseSchemaCompat = ApiKeyCreateResponseSchema[_Actions]
+
+
+class ApiKeySchema(typing.Generic[_ActionT], typing.TypedDict):
     """
     Response schema for an [API key](https://typesense.org/docs/26.0/api/api-keys.html#retrieve-an-api-key).
 
     Attributes:
-        actions (list[str]): The actions allowed for this key.
+        actions (list[_ActionT]): The actions allowed for this key.
         collections (list[str]): The collections this key has access to.
         description (str): The description for this key.
         id (int): The ID of the key.
         value_prefix (str): The value prefix of the key.
         expires_at (int): The time in UNIX timestamp format when the key
     """
 
-    actions: typing.List[_Actions]
+    actions: typing.List[_ActionT]
     collections: typing.List[str]
     description: str
     id: int
     value_prefix: str
     expires_at: int
 
 
-class ApiKeyRetrieveSchema(typing.TypedDict):
+ApiKeySchemaCompat = ApiKeySchema[_Actions]
+
+
+class ApiKeyRetrieveSchema(typing.Generic[_ActionT], typing.TypedDict):
     """
     Response schema for retrieving [API keys](https://typesense.org/docs/26.0/api/api-keys.html#list-all-keys).
 
     Attributes:
-        keys (list[ApiKeySchema]): The list of keys.
+        keys (list[ApiKeySchema[_ActionT]]): The list of keys.
     """
 
-    keys: typing.List[ApiKeySchema]
+    keys: typing.List[ApiKeySchema[_ActionT]]
+
+
+ApiKeyRetrieveSchemaCompat = ApiKeyRetrieveSchema[_Actions]
 
 
 class ApiKeyDeleteSchema(typing.TypedDict):

src/typesense/types/override.py

@@ -8,6 +8,12 @@
     import typing_extensions as typing
 
 
+_OverrideRuleT = typing.TypeVar(
+    "_OverrideRuleT",
+    bound=typing.Union["OverrideQueryRuleSchema", "OverrideFilterSchema"],
+)
+
+
 class OverrideQueryRuleSchema(typing.TypedDict):
     """
     The schema for the rule field in the Overrides.upsert method.
@@ -51,12 +57,12 @@ class IncludesSchema(typing.TypedDict):
     position: int
 
 
-class OverrideCreateSchema(typing.TypedDict):
+class OverrideCreateSchema(typing.Generic[_OverrideRuleT], typing.TypedDict):
     """
     The schema for the request of the Overrides.upsert method.
 
     Attributes:
-        rule (OverrideQueryRuleSchema | OverrideFilterSchema): The rule.
+        rule (_OverrideRuleT): The rule.
         sort_by (str): The sort by string.
         filter_by (str): The filter by string.
         excludes (list[str]): The excludes list.
@@ -69,7 +75,7 @@ class OverrideCreateSchema(typing.TypedDict):
         stop_processing (bool): Whether to stop processing.
     """
 
-    rule: typing.Union[OverrideQueryRuleSchema, OverrideFilterSchema]
+    rule: _OverrideRuleT
     sort_by: typing.NotRequired[str]
     filter_by: typing.NotRequired[str]
     excludes: typing.NotRequired[typing.List[str]]
@@ -82,12 +88,22 @@ class OverrideCreateSchema(typing.TypedDict):
     stop_processing: typing.NotRequired[bool]
 
 
-class OverrideSchema(OverrideCreateSchema):
+OverrideCreateSchemaCompat = OverrideCreateSchema[
+    typing.Union[OverrideQueryRuleSchema, OverrideFilterSchema]
+]
+
+
+class OverrideSchema(OverrideCreateSchema[_OverrideRuleT]):
     """The schema for the response of the Overrides.upsert method."""
 
     id: str
 
 
+OverrideSchemaCompat = OverrideSchema[
+    typing.Union[OverrideQueryRuleSchema, OverrideFilterSchema]
+]
+
+
 class OverrideDeleteSchema(typing.TypedDict):
     """The schema for the response of the Overrides.delete method."""
 
@@ -97,4 +113,4 @@ class OverrideDeleteSchema(typing.TypedDict):
 class OverrideRetrieveSchema(typing.TypedDict):
     """The schema for the response of the Overrides.retrieve method."""
 
-    overrides: typing.List[OverrideSchema]
+    overrides: typing.List[OverrideSchema[typing.Union[OverrideQueryRuleSchema, OverrideFilterSchema]]]