Merge pull request #1007 from CAPITAINMARVEL/union_doc_fix

Fix UnionDoc type hint missing in init_beanie and on FindInterface

beanie/odm/interfaces/find.py

@@ -27,10 +27,11 @@
 
 if TYPE_CHECKING:
     from beanie.odm.documents import Document
+    from beanie.odm.union_doc import UnionDoc
     from beanie.odm.views import View
 
 DocumentProjectionType = TypeVar("DocumentProjectionType", bound=BaseModel)
-FindType = TypeVar("FindType", bound=Union["Document", "View"])
+FindType = TypeVar("FindType", bound=Union["Document", "UnionDoc", "View"])
 
 
 class FindInterface:

beanie/odm/union_doc.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, Dict, Optional, Type
+from typing import ClassVar, Dict, Optional, Type, TypeVar
 
 from beanie.exceptions import UnionDocNotInited
 from beanie.odm.interfaces.aggregate import AggregateInterface
@@ -7,6 +7,8 @@
 from beanie.odm.interfaces.getters import OtherGettersInterface
 from beanie.odm.settings.union_doc import UnionDocSettings
 
+UnionDocType = TypeVar("UnionDocType", bound="UnionDoc")
+
 
 class UnionDoc(
     FindInterface,

beanie/odm/utils/init.py

@@ -50,7 +50,7 @@
 from beanie.odm.settings.document import DocumentSettings, IndexModelField
 from beanie.odm.settings.union_doc import UnionDocSettings
 from beanie.odm.settings.view import ViewSettings
-from beanie.odm.union_doc import UnionDoc
+from beanie.odm.union_doc import UnionDoc, UnionDocType
 from beanie.odm.views import View
 
 
@@ -65,7 +65,9 @@ def __init__(
         database: AsyncIOMotorDatabase = None,
         connection_string: Optional[str] = None,
         document_models: Optional[
-            List[Union[Type["DocType"], Type["View"], str]]
+            List[
+                Union[Type["DocType"], Type["UnionDocType"], Type["View"], str]
+            ]
         ] = None,
         allow_index_dropping: bool = False,
         recreate_views: bool = False,
@@ -76,7 +78,7 @@ def __init__(
 
         :param database: AsyncIOMotorDatabase - motor database instance
         :param connection_string: str - MongoDB connection string
-        :param document_models: List[Union[Type[DocType], str]] - model classes
+        :param document_models: List[Union[Type[DocType], Type[UnionDocType], str]] - model classes
         or strings with dot separated paths
         :param allow_index_dropping: bool - if index dropping is allowed.
         Default False
@@ -117,7 +119,9 @@ def __init__(
             ModelType.View: 2,
         }
 
-        self.document_models: List[Union[Type[DocType], Type[View]]] = [
+        self.document_models: List[
+            Union[Type[DocType], Type[UnionDocType], Type[View]]
+        ] = [
             self.get_model(model) if isinstance(model, str) else model
             for model in document_models
         ]
@@ -757,7 +761,7 @@ async def init_beanie(
     database: AsyncIOMotorDatabase = None,
     connection_string: Optional[str] = None,
     document_models: Optional[
-        List[Union[Type[Document], Type["View"], str]]
+        List[Union[Type[Document], Type[UnionDoc], Type["View"], str]]
     ] = None,
     allow_index_dropping: bool = False,
     recreate_views: bool = False,
@@ -768,7 +772,7 @@ async def init_beanie(
 
     :param database: AsyncIOMotorDatabase - motor database instance
     :param connection_string: str - MongoDB connection string
-    :param document_models: List[Union[Type[DocType], str]] - model classes
+    :param document_models: List[Union[Type[DocType], Type[UnionDocType], str]] - model classes
     or strings with dot separated paths
     :param allow_index_dropping: bool - if index dropping is allowed.
     Default False