feat add overloads for field types to improve type safety

ormar/fields/model_fields.py

@@ -3,7 +3,7 @@
 import uuid
 from enum import Enum as E
 from enum import EnumMeta
-from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union, overload
+from typing import Any, Optional, Type
 
 import pydantic
 import sqlalchemy
@@ -14,11 +14,6 @@
 from ormar.fields.base import BaseField  # noqa I101
 from ormar.fields.sqlalchemy_encrypted import EncryptBackends
 
-try:
-    from typing import Literal  # type: ignore
-except ImportError:  # pragma: no cover
-    from typing_extensions import Literal  # type: ignore
-
 try:
     from typing import Self  # type: ignore
 except ImportError:  # pragma: no cover
@@ -327,33 +322,26 @@ def get_column_type(cls, **kwargs: Any) -> Any:
         return sqlalchemy.Float()
 
 
-if TYPE_CHECKING:  # pragma: nocover
-
-    def Boolean(**kwargs: Any) -> bool:
-        pass
+class Boolean(ModelFieldFactory, int):
+    """
+    Boolean field factory that construct Field classes and populated their values.
+    """
 
-else:
+    _type = bool
+    _sample = True
 
-    class Boolean(ModelFieldFactory, int):
-        """
-        Boolean field factory that construct Field classes and populated their values.
+    @classmethod
+    def get_column_type(cls, **kwargs: Any) -> Any:
         """
+        Return proper type of db column for given field type.
+        Accepts required and optional parameters that each column type accepts.
 
-        _type = bool
-        _sample = True
-
-        @classmethod
-        def get_column_type(cls, **kwargs: Any) -> Any:
-            """
-            Return proper type of db column for given field type.
-            Accepts required and optional parameters that each column type accepts.
-
-            :param kwargs: key, value pairs of sqlalchemy options
-            :type kwargs: Any
-            :return: initialized column with proper options
-            :rtype: sqlalchemy Column
-            """
-            return sqlalchemy.Boolean()
+        :param kwargs: key, value pairs of sqlalchemy options
+        :type kwargs: Any
+        :return: initialized column with proper options
+        :rtype: sqlalchemy Column
+        """
+        return sqlalchemy.Boolean()
 
 
 class DateTime(ModelFieldFactory, datetime.datetime):
@@ -470,81 +458,53 @@ def get_column_type(cls, **kwargs: Any) -> Any:
         return sqlalchemy.JSON(none_as_null=kwargs.get("sql_nullable", False))
 
 
-if TYPE_CHECKING:  # pragma: nocover # noqa: C901
-
-    @overload
-    def LargeBinary(  # type: ignore
-        max_length: int, *, represent_as_base64_str: Literal[True], **kwargs: Any
-    ) -> str: ...
-
-    @overload
-    def LargeBinary(  # type: ignore
-        max_length: int, *, represent_as_base64_str: Literal[False], **kwargs: Any
-    ) -> bytes: ...
-
-    @overload
-    def LargeBinary(
-        max_length: int, represent_as_base64_str: Literal[False] = ..., **kwargs: Any
-    ) -> bytes: ...
-
-    def LargeBinary(
-        max_length: int, represent_as_base64_str: bool = False, **kwargs: Any
-    ) -> Union[str, bytes]:
-        pass
-
-else:
-
-    class LargeBinary(ModelFieldFactory, bytes):
-        """
-        LargeBinary field factory that construct Field classes
-        and populated their values.
-        """
-
-        _type = bytes
-        _sample = "bytes"
-
-        def __new__(  # type: ignore # noqa CFQ002
-            cls,
-            *,
-            max_length: int,
-            represent_as_base64_str: bool = False,
-            **kwargs: Any
-        ) -> Self:  # type: ignore
-            kwargs = {
-                **kwargs,
-                **{
-                    k: v
-                    for k, v in locals().items()
-                    if k not in ["cls", "__class__", "kwargs"]
-                },
-            }
-            return super().__new__(cls, **kwargs)
-
-        @classmethod
-        def get_column_type(cls, **kwargs: Any) -> Any:
-            """
-            Return proper type of db column for given field type.
-            Accepts required and optional parameters that each column type accepts.
-
-            :param kwargs: key, value pairs of sqlalchemy options
-            :type kwargs: Any
-            :return: initialized column with proper options
-            :rtype: sqlalchemy Column
-            """
-            return sqlalchemy.LargeBinary(length=kwargs.get("max_length"))
-
-        @classmethod
-        def validate(cls, **kwargs: Any) -> None:
-            """
-            Used to validate if all required parameters on a given field type are set.
-            :param kwargs: all params passed during construction
-            :type kwargs: Any
-            """
-            max_length = kwargs.get("max_length", None)
-            if max_length <= 0:
-                raise ModelDefinitionError(
-                    "Parameter max_length is required for field LargeBinary"
-                )
+class LargeBinary(ModelFieldFactory, bytes):
+    """
+    LargeBinary field factory that construct Field classes
+    and populated their values.
+    """
+
+    _type = bytes
+    _sample = "bytes"
+
+    def __new__(  # type: ignore # noqa CFQ002
+        cls, *, max_length: int, represent_as_base64_str: bool = False, **kwargs: Any
+    ) -> Self:  # type: ignore
+        kwargs = {
+            **kwargs,
+            **{
+                k: v
+                for k, v in locals().items()
+                if k not in ["cls", "__class__", "kwargs"]
+            },
+        }
+        return super().__new__(cls, **kwargs)
+
+    @classmethod
+    def get_column_type(cls, **kwargs: Any) -> Any:
+        """
+        Return proper type of db column for given field type.
+        Accepts required and optional parameters that each column type accepts.
+
+        :param kwargs: key, value pairs of sqlalchemy options
+        :type kwargs: Any
+        :return: initialized column with proper options
+        :rtype: sqlalchemy Column
+        """
+        return sqlalchemy.LargeBinary(length=kwargs.get("max_length"))
+
+    @classmethod
+    def validate(cls, **kwargs: Any) -> None:
+        """
+        Used to validate if all required parameters on a given field type are set.
+        :param kwargs: all params passed during construction
+        :type kwargs: Any
+        """
+        max_length = kwargs.get("max_length", None)
+        if max_length <= 0:
+            raise ModelDefinitionError(
+                "Parameter max_length is required for field LargeBinary"
+            )
 
 
 class BigInteger(Integer, int):
@@ -753,42 +713,34 @@ def get_column_type(cls, **kwargs: Any) -> Any:
         return sqlalchemy_uuid.UUID(uuid_format=uuid_format)
 
 
-if TYPE_CHECKING:  # pragma: nocover
-    T = TypeVar("T", bound=E)
-
-    def Enum(enum_class: Type[T], **kwargs: Any) -> T:
-        pass
-
-else:
-
-    class Enum(ModelFieldFactory):
-        """
-        Enum field factory that construct Field classes and populated their values.
-        """
+class Enum(ModelFieldFactory):
+    """
+    Enum field factory that construct Field classes and populated their values.
+    """
 
-        _type = E
-        _sample = None
+    _type = E
+    _sample = None
 
-        def __new__(  # type: ignore # noqa CFQ002
-            cls, *, enum_class: Type[E], **kwargs: Any
-        ) -> Self:
-            kwargs = {
-                **kwargs,
-                **{
-                    k: v
-                    for k, v in locals().items()
-                    if k not in ["cls", "__class__", "kwargs"]
-                },
-            }
-            return super().__new__(cls, **kwargs)
+    def __new__(  # type: ignore # noqa CFQ002
+        cls, *, enum_class: Type[E], **kwargs: Any
+    ) -> Self:
+        kwargs = {
+            **kwargs,
+            **{
+                k: v
+                for k, v in locals().items()
+                if k not in ["cls", "__class__", "kwargs"]
+            },
+        }
+        return super().__new__(cls, **kwargs)
 
-        @classmethod
-        def validate(cls, **kwargs: Any) -> None:
-            enum_class = kwargs.get("enum_class")
-            if enum_class is None or not isinstance(enum_class, EnumMeta):
-                raise ModelDefinitionError("Enum Field choices must be EnumType")
+    @classmethod
+    def validate(cls, **kwargs: Any) -> None:
+        enum_class = kwargs.get("enum_class")
+        if enum_class is None or not isinstance(enum_class, EnumMeta):
+            raise ModelDefinitionError("Enum Field choices must be EnumType")
 
-        @classmethod
-        def get_column_type(cls, **kwargs: Any) -> Any:
-            enum_cls = kwargs.get("enum_class")
-            return sqlalchemy.Enum(enum_cls)
+    @classmethod
+    def get_column_type(cls, **kwargs: Any) -> Any:
+        enum_cls = kwargs.get("enum_class")
+        return sqlalchemy.Enum(enum_cls)

ormar/fields/model_fields.pyi

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from datetime import date, datetime, time
+from decimal import Decimal as DecimalType
+from enum import Enum as EnumBase
+from typing import Any, Literal, Type, TypeVar, Union, overload
+from uuid import UUID as UuidType
+
+T = TypeVar("T", bound=EnumBase)
+
+@overload
+def Boolean(*, nullable: Literal[False] = False, **kwargs: Any) -> bool: ...
+@overload
+def Boolean(*, nullable: Literal[True], **kwargs: Any) -> bool | None: ...
+@overload
+def String(
+    *, max_length: int, nullable: Literal[False] = False, **kwargs: Any
+) -> str: ...
+@overload
+def String(
+    *, max_length: int, nullable: Literal[True], **kwargs: Any
+) -> str | None: ...
+@overload
+def Integer(*, nullable: Literal[False] = False, **kwargs: Any) -> int: ...
+@overload
+def Integer(*, nullable: Literal[True], **kwargs: Any) -> int | None: ...
+@overload
+def Text(*, nullable: Literal[False] = False, **kwargs: Any) -> str: ...
+@overload
+def Text(*, nullable: Literal[True], **kwargs: Any) -> str | None: ...
+@overload
+def Float(*, nullable: Literal[False] = False, **kwargs: Any) -> float: ...
+@overload
+def Float(*, nullable: Literal[True], **kwargs: Any) -> float | None: ...
+@overload
+def DateTime(*, nullable: Literal[False] = False, **kwargs: Any) -> datetime: ...
+@overload
+def DateTime(*, nullable: Literal[True], **kwargs: Any) -> datetime | None: ...
+@overload
+def Date(*, nullable: Literal[False] = False, **kwargs: Any) -> date: ...
+@overload
+def Date(*, nullable: Literal[True], **kwargs: Any) -> date | None: ...
+@overload
+def Time(*, nullable: Literal[False] = False, **kwargs: Any) -> time: ...
+@overload
+def Time(*, nullable: Literal[True], **kwargs: Any) -> time | None: ...
+@overload
+def JSON(*, nullable: Literal[False] = False, **kwargs: Any) -> Any: ...
+@overload
+def JSON(*, nullable: Literal[True], **kwargs: Any) -> Any | None: ...
+@overload
+def BigInteger(*, nullable: Literal[False] = False, **kwargs: Any) -> int: ...
+@overload
+def BigInteger(*, nullable: Literal[True], **kwargs: Any) -> int | None: ...
+@overload
+def SmallInteger(*, nullable: Literal[False] = False, **kwargs: Any) -> int: ...
+@overload
+def SmallInteger(*, nullable: Literal[True], **kwargs: Any) -> int | None: ...
+@overload
+def Decimal(
+    *,
+    max_digits: int,
+    decimal_places: int,
+    nullable: Literal[False] = False,
+    **kwargs: Any,
+) -> DecimalType: ...
+@overload
+def Decimal(
+    *, max_digits: int, decimal_places: int, nullable: Literal[True], **kwargs: Any
+) -> DecimalType | None: ...
+@overload
+def UUID(*, nullable: Literal[False] = False, **kwargs: Any) -> UuidType: ...
+@overload
+def UUID(*, nullable: Literal[True], **kwargs: Any) -> UuidType | None: ...
+@overload
+def LargeBinary(
+    max_length: int, *, represent_as_base64_str: Literal[True], **kwargs: Any
+) -> str: ...
+@overload
+def LargeBinary(
+    max_length: int, *, represent_as_base64_str: Literal[False], **kwargs: Any
+) -> bytes: ...
+@overload
+def LargeBinary(
+    max_length: int, represent_as_base64_str: Literal[False] = ..., **kwargs: Any
+) -> bytes: ...
+@overload
+def LargeBinary(
+    max_length: int, represent_as_base64_str: bool = False, **kwargs: Any
+) -> Union[str, bytes]: ...
+@overload
+def Enum(
+    enum_class: Type[T],
+    *,
+    nullable: Literal[False] = False,
+    **kwargs: Any,
+) -> T: ...
+@overload
+def Enum(
+    enum_class: Type[T],
+    *,
+    nullable: Literal[True],
+    **kwargs: Any,
+) -> T | None: ...

tests/test_encryption/test_encrypted_columns.py

@@ -4,7 +4,7 @@
 import decimal
 import hashlib
 import uuid
-from typing import Any
+from typing import Any, Optional
 
 import ormar
 import pytest
@@ -51,7 +51,7 @@ class Author(ormar.Model):
     test_text: str = ormar.Text(default="", **default_fernet)
     test_bool: bool = ormar.Boolean(nullable=False, **default_fernet)
     test_float: float = ormar.Float(**default_fernet)
-    test_float2: float = ormar.Float(nullable=True, **default_fernet)
+    test_float2: Optional[float] = ormar.Float(nullable=True, **default_fernet)
     test_datetime = ormar.DateTime(default=datetime.datetime.now, **default_fernet)
     test_date = ormar.Date(default=datetime.date.today, **default_fernet)
     test_time = ormar.Time(default=datetime.time, **default_fernet)