Go/missing typing (sbdchd#32)

Add missing typing for `register_connection` and `EnumField` according to [API](http://docs.mongoengine.org/apireference.html#mongoengine.register_connection) for `register_connection` and [the source](https://sourcegraph.com/github.com/MongoEngine/mongoengine/-/blob/mongoengine/fields.py#L1610) for `EnumField`

tests/test_mongoengine.py

@@ -2,7 +2,8 @@
 
 import datetime
 import types
-from typing import Any, KeysView, Type, TypeVar, cast
+from enum import Enum
+from typing import Any, KeysView, Type, TypeVar, Union, cast
 
 import mongoengine
 import pymongo
@@ -50,6 +51,12 @@ class Recipe(Document):
     description = fields.StringField()
 
 
+class Font(Enum):
+    Helvetica = ("Helvetica",)
+    Arial = ("Arial",)
+    Times = "Times New Roman"
+
+
 class Post(Document):
     meta = {
         "collection": "posts",
@@ -88,6 +95,13 @@ def dead_posts(cls) -> QuerySet[Post]:
         help_text=("Map tag names to descriptions"),
     )
 
+    font = fields.EnumField(Font)
+    font_required = fields.EnumField(Font, required=True)
+    font_default = fields.EnumField(Font, default=Font.Helvetica)
+    font_required_default = fields.EnumField(
+        Font, required=True, default=Font.Helvetica
+    )
+
     def set_hidden(self, hidden: bool) -> None:
         self.hidden = hidden
         self.save()
@@ -144,6 +158,21 @@ def main() -> None:
     first_post.tags.values()
     first_post.errors
     first_post.results
+    first_post.font
+
+    def log_optional_font(f: Union[Font, None]) -> None:
+        print(f)
+
+    log_optional_font(first_post.font)
+
+    def log_required_font(f: Font) -> None:
+        print(f)
+
+    log_required_font(first_post.font_required)
+    first_post.font_required_default = None
+
+    log_required_font(first_post.font_default)
+    first_post.font_default = None
 
     p = Post()
     p.validate()

typings/mongoengine/__init__.pyi

@@ -6,10 +6,21 @@ from mongoengine.document import Document, DynamicDocument, EmbeddedDocument
 from mongoengine.errors import DoesNotExist, NotUniqueError, ValidationError
 from mongoengine.queryset.queryset import QuerySet
 from mongoengine.queryset.visitor import Q
-from pymongo import MongoClient
+from pymongo import MongoClient, ReadPreference
 
 def connect(name: str, alias: str = ..., host: Optional[str] = ...) -> MongoClient: ...
-def register_connection(host: str, alias: str) -> None: ...
+def register_connection(
+    alias: str,
+    db: Optional[str] = ...,
+    name: Optional[str] = ...,
+    host: Optional[str] = ...,
+    port: Optional[int] = ...,
+    read_preference: ReadPreference = ...,
+    username: Optional[str] = ...,
+    password: Optional[str] = ...,
+    authentication_source: Optional[str] = ...,
+    authentication_mechanism: Optional[str] = ...,
+) -> None: ...
 
 __all__ = [
     "Q",

typings/mongoengine/fields.pyi

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from datetime import datetime
 from decimal import Decimal
+from enum import Enum
 from typing import (
     Any,
     Callable,
@@ -1068,3 +1069,77 @@ class ReferenceField(BaseField):
         blank: bool = ...,
     ) -> None: ...
     def __getitem__(self, arg: Any) -> Any: ...
+
+_T_ENUM = TypeVar("_T_ENUM", bound=Enum)
+
+class EnumField(Generic[_ST, _GT], BaseField):
+    @overload
+    def __new__(
+        cls,
+        enum: Type[_T_ENUM],
+        *,
+        db_field: str = ...,
+        name: Optional[str] = ...,
+        required: Literal[False] = ...,
+        default: None = ...,
+        unique: bool = ...,
+        unique_with: Union[str, Iterable[str]] = ...,
+        primary_key: Literal[False] = ...,
+        choices: Optional[Iterable[_T_ENUM]] = ...,
+        null: bool = ...,
+        verbose_name: Optional[str] = ...,
+        help_text: Optional[str] = ...,
+    ) -> EnumField[Optional[_T_ENUM], Optional[_T_ENUM]]: ...
+    @overload
+    def __new__(
+        cls,
+        enum: Type[_T_ENUM],
+        *,
+        db_field: str = ...,
+        name: Optional[str] = ...,
+        required: Literal[False] = ...,
+        default: Union[_T_ENUM, Callable[[], _T_ENUM]],
+        unique: bool = ...,
+        unique_with: Union[str, Iterable[str]] = ...,
+        primary_key: Literal[False] = ...,
+        choices: Optional[Iterable[_T_ENUM]] = ...,
+        null: bool = ...,
+        verbose_name: Optional[str] = ...,
+        help_text: Optional[str] = ...,
+    ) -> EnumField[Optional[_T_ENUM], _T_ENUM]: ...
+    @overload
+    def __new__(
+        cls,
+        enum: Type[_T_ENUM],
+        *,
+        db_field: str = ...,
+        name: Optional[str] = ...,
+        required: Literal[True],
+        default: None = ...,
+        unique: bool = ...,
+        unique_with: Union[str, Iterable[str]] = ...,
+        primary_key: Literal[False] = ...,
+        choices: Optional[Iterable[_T_ENUM]] = ...,
+        null: bool = ...,
+        verbose_name: Optional[str] = ...,
+        help_text: Optional[str] = ...,
+    ) -> EnumField[_T_ENUM, _T_ENUM]: ...
+    @overload
+    def __new__(
+        cls,
+        enum: Type[_T_ENUM],
+        *,
+        db_field: str = ...,
+        name: Optional[str] = ...,
+        required: Literal[True],
+        default: Union[_T_ENUM, Callable[[], _T_ENUM]],
+        unique: bool = ...,
+        unique_with: Union[str, Iterable[str]] = ...,
+        primary_key: Literal[False] = ...,
+        choices: Optional[Iterable[_T_ENUM]] = ...,
+        null: bool = ...,
+        verbose_name: Optional[str] = ...,
+        help_text: Optional[str] = ...,
+    ) -> EnumField[Optional[_T_ENUM], _T_ENUM]: ...
+    def __set__(self, instance: Any, value: _ST) -> None: ...
+    def __get__(self, instance: Any, owner: Any) -> _GT: ...