feat add overloads for field types to improve type safety

ormar/fields/model_fields.pyi

@@ -0,0 +1,84 @@
+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, overload
+from uuid import UUID as UuidType
+
+T = TypeVar("T", bound=EnumBase)
+
+def Boolean(**kwargs: Any) -> bool: ...
+@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: ...
+def Enum(enum_class: Type[T], **kwargs: Any) -> T: ...

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)

tests/test_exclude_include_dict/test_dumping_model_to_dict.py

@@ -21,7 +21,7 @@ class User(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     email: str = ormar.String(max_length=255, nullable=False)
-    password: str = ormar.String(max_length=255, nullable=True)
+    password: Optional[str] = ormar.String(max_length=255, nullable=True)
     first_name: str = ormar.String(max_length=255, nullable=False)
     roles: List[Role] = ormar.ManyToMany(Role)
 

tests/test_exclude_include_dict/test_excludable_items.py

@@ -14,7 +14,7 @@ class NickNames(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     name: str = ormar.String(max_length=100, nullable=False, name="hq_name")
-    is_lame: bool = ormar.Boolean(nullable=True)
+    is_lame: Optional[bool] = ormar.Boolean(nullable=True)
 
 
 class NicksHq(ormar.Model):
@@ -34,7 +34,7 @@ class Company(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     name: str = ormar.String(max_length=100, nullable=False, name="company_name")
-    founded: int = ormar.Integer(nullable=True)
+    founded: Optional[int] = ormar.Integer(nullable=True)
     hq: HQ = ormar.ForeignKey(HQ)
 
 
@@ -44,10 +44,10 @@ class Car(ormar.Model):
     id: int = ormar.Integer(primary_key=True)
     manufacturer: Optional[Company] = ormar.ForeignKey(Company)
     name: str = ormar.String(max_length=100)
-    year: int = ormar.Integer(nullable=True)
-    gearbox_type: str = ormar.String(max_length=20, nullable=True)
-    gears: int = ormar.Integer(nullable=True)
-    aircon_type: str = ormar.String(max_length=20, nullable=True)
+    year: Optional[int] = ormar.Integer(nullable=True)
+    gearbox_type: Optional[str] = ormar.String(max_length=20, nullable=True)
+    gears: Optional[int] = ormar.Integer(nullable=True)
+    aircon_type: Optional[str] = ormar.String(max_length=20, nullable=True)
 
 
 create_test_database = init_tests(base_ormar_config)

tests/test_exclude_include_dict/test_excluding_fields_in_fastapi.py

@@ -1,6 +1,7 @@
 import datetime
 import random
 import string
+from typing import Optional
 
 import ormar
 import pydantic
@@ -64,10 +65,10 @@ class User(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     email: str = ormar.String(max_length=255)
-    password: str = ormar.String(max_length=255, nullable=True)
+    password: Optional[str] = ormar.String(max_length=255, nullable=True)
     first_name: str = ormar.String(max_length=255)
     last_name: str = ormar.String(max_length=255)
-    category: str = ormar.String(max_length=255, nullable=True)
+    category: Optional[str] = ormar.String(max_length=255, nullable=True)
 
 
 class User2(ormar.Model):
@@ -78,7 +79,7 @@ class User2(ormar.Model):
     password: str = ormar.String(max_length=255)
     first_name: str = ormar.String(max_length=255)
     last_name: str = ormar.String(max_length=255)
-    category: str = ormar.String(max_length=255, nullable=True)
+    category: Optional[str] = ormar.String(max_length=255, nullable=True)
     timestamp: datetime.datetime = pydantic.Field(default_factory=datetime.datetime.now)
 
 

tests/test_exclude_include_dict/test_excluding_fields_with_default.py

@@ -19,7 +19,7 @@ class Album(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     name: str = ormar.String(max_length=100)
-    is_best_seller: bool = ormar.Boolean(default=False, nullable=True)
+    is_best_seller: Optional[bool] = ormar.Boolean(default=False, nullable=True)
 
 
 class Track(ormar.Model):
@@ -29,7 +29,7 @@ class Track(ormar.Model):
     album: Optional[Album] = ormar.ForeignKey(Album)
     title: str = ormar.String(max_length=100)
     position: int = ormar.Integer(default=get_position)
-    play_count: int = ormar.Integer(nullable=True, default=0)
+    play_count: Optional[int] = ormar.Integer(nullable=True, default=0)
 
 
 create_test_database = init_tests(base_ormar_config)

tests/test_exclude_include_dict/test_excluding_subset_of_columns.py

@@ -15,7 +15,7 @@ class Company(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     name: str = ormar.String(max_length=100, nullable=False)
-    founded: int = ormar.Integer(nullable=True)
+    founded: Optional[int] = ormar.Integer(nullable=True)
 
 
 class Car(ormar.Model):
@@ -24,10 +24,10 @@ class Car(ormar.Model):
     id: int = ormar.Integer(primary_key=True)
     manufacturer: Optional[Company] = ormar.ForeignKey(Company)
     name: str = ormar.String(max_length=100)
-    year: int = ormar.Integer(nullable=True)
-    gearbox_type: str = ormar.String(max_length=20, nullable=True)
-    gears: int = ormar.Integer(nullable=True, name="gears_number")
-    aircon_type: str = ormar.String(max_length=20, nullable=True)
+    year: Optional[int] = ormar.Integer(nullable=True)
+    gearbox_type: Optional[str] = ormar.String(max_length=20, nullable=True)
+    gears: Optional[int] = ormar.Integer(nullable=True, name="gears_number")
+    aircon_type: Optional[str] = ormar.String(max_length=20, nullable=True)
 
 
 create_test_database = init_tests(base_ormar_config)

tests/test_exclude_include_dict/test_pydantic_dict_params.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import ormar
 import pytest
@@ -13,7 +13,7 @@ class Category(ormar.Model):
     ormar_config = base_ormar_config.copy(tablename="categories")
 
     id: int = ormar.Integer(primary_key=True)
-    name: str = ormar.String(max_length=100, default="Test", nullable=True)
+    name: Optional[str] = ormar.String(max_length=100, default="Test", nullable=True)
     visibility: bool = ormar.Boolean(default=True)
 
 

tests/test_fastapi/test_json_field_fastapi.py

@@ -1,6 +1,6 @@
 # type: ignore
 import uuid
-from typing import List
+from typing import List, Optional
 
 import ormar
 import pydantic
@@ -79,7 +79,7 @@ class Thing2(ormar.Model):
 
         id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
         name: str = ormar.Text(default="")
-        js: pydantic.Json = ormar.JSON(nullable=True)
+        js: Optional[pydantic.Json] = ormar.JSON(nullable=True)
 
     Thing2()
 

tests/test_fastapi/test_recursion_error.py

@@ -1,6 +1,6 @@
 import uuid
 from datetime import datetime
-from typing import List
+from typing import List, Optional
 
 import ormar
 import pytest
@@ -27,7 +27,7 @@ class User(ormar.Model):
     username: str = ormar.String(unique=True, max_length=100)
     password: str = ormar.String(unique=True, max_length=100)
     verified: bool = ormar.Boolean(default=False)
-    verify_key: str = ormar.String(unique=True, max_length=100, nullable=True)
+    verify_key: Optional[str] = ormar.String(unique=True, max_length=100, nullable=True)
     created_at: datetime = ormar.DateTime(default=datetime.now())
 
     ormar_config = base_ormar_config.copy(tablename="users")
@@ -65,7 +65,7 @@ class QuizInput(BaseModel):
 class Quiz(ormar.Model):
     id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
     title: str = ormar.String(max_length=100)
-    description: str = ormar.String(max_length=300, nullable=True)
+    description: Optional[str] = ormar.String(max_length=300, nullable=True)
     created_at: datetime = ormar.DateTime(default=datetime.now())
     updated_at: datetime = ormar.DateTime(default=datetime.now())
     user_id: uuid.UUID = ormar.UUID(foreign_key=User.id)

tests/test_fastapi/test_relations_with_nested_defaults.py

@@ -36,7 +36,7 @@ class Book(ormar.Model):
     id: int = ormar.Integer(primary_key=True)
     author: Optional[Author] = ormar.ForeignKey(Author)
     title: str = ormar.String(max_length=100)
-    year: int = ormar.Integer(nullable=True)
+    year: Optional[int] = ormar.Integer(nullable=True)
 
 
 create_test_database = init_tests(base_ormar_config)

tests/test_fastapi/test_wekref_exclusion.py

@@ -28,7 +28,7 @@ class Thing(ormar.Model):
 
     id: UUID = ormar.UUID(primary_key=True, default=uuid4)
     name: str = ormar.Text(default="")
-    js: pydantic.Json = ormar.JSON(nullable=True)
+    js: Optional[pydantic.Json] = ormar.JSON(nullable=True)
     other_thing: Optional[OtherThing] = ormar.ForeignKey(OtherThing, nullable=True)
 
 

tests/test_inheritance_and_pydantic_generation/test_inherited_class_is_not_abstract_by_default.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Optional
 
 import ormar
 import pytest
@@ -17,8 +18,10 @@ class TableBase(ormar.Model):
     created_at: datetime.datetime = ormar.DateTime(
         timezone=True, default=datetime.datetime.now
     )
-    last_modified_by: str = ormar.String(max_length=20, nullable=True)
-    last_modified_at: datetime.datetime = ormar.DateTime(timezone=True, nullable=True)
+    last_modified_by: Optional[str] = ormar.String(max_length=20, nullable=True)
+    last_modified_at: Optional[datetime.datetime] = ormar.DateTime(
+        timezone=True, nullable=True
+    )
 
 
 class NationBase(ormar.Model):

tests/test_model_definition/test_aliases.py

@@ -15,7 +15,7 @@ class Child(ormar.Model):
     id: int = ormar.Integer(name="child_id", primary_key=True)
     first_name: str = ormar.String(name="fname", max_length=100)
     last_name: str = ormar.String(name="lname", max_length=100)
-    born_year: int = ormar.Integer(name="year_born", nullable=True)
+    born_year: Optional[int] = ormar.Integer(name="year_born", nullable=True)
 
 
 class Artist(ormar.Model):

tests/test_model_definition/test_columns.py

@@ -1,5 +1,6 @@
 import datetime
 from enum import Enum
+from typing import Optional
 
 import ormar
 import pydantic
@@ -29,8 +30,8 @@ class Example(ormar.Model):
     created: datetime.datetime = ormar.DateTime(default=datetime.datetime.now)
     created_day: datetime.date = ormar.Date(default=datetime.date.today)
     created_time: datetime.time = ormar.Time(default=time)
-    description: str = ormar.Text(nullable=True)
-    value: float = ormar.Float(nullable=True)
+    description: Optional[str] = ormar.Text(nullable=True)
+    value: Optional[float] = ormar.Float(nullable=True)
     data: pydantic.Json = ormar.JSON(default={})
     size: MyEnum = ormar.Enum(enum_class=MyEnum, default=MyEnum.SMALL)
 

tests/test_model_definition/test_model_construct.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import ormar
 import pytest
@@ -33,7 +33,7 @@ class Company(ormar.Model):
 
     id: int = ormar.Integer(primary_key=True)
     name: str = ormar.String(max_length=100, nullable=False, name="company_name")
-    founded: int = ormar.Integer(nullable=True)
+    founded: Optional[int] = ormar.Integer(nullable=True)
     hq: HQ = ormar.ForeignKey(HQ)
 
 

tests/test_model_definition/test_models.py

@@ -4,6 +4,7 @@
 import os
 import uuid
 from enum import Enum
+from typing import Optional
 
 import ormar
 import pydantic
@@ -126,7 +127,9 @@ class NullableCountry(ormar.Model):
     ormar_config = base_ormar_config.copy(tablename="country2")
 
     id: int = ormar.Integer(primary_key=True)
-    name: CountryNameEnum = ormar.Enum(enum_class=CountryNameEnum, nullable=True)
+    name: Optional[CountryNameEnum] = ormar.Enum(
+        enum_class=CountryNameEnum, nullable=True
+    )
 
 
 class NotNullableCountry(ormar.Model):

tests/test_model_definition/test_pk_field_is_always_not_null.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import ormar
 
 from tests.lifespan import init_tests
@@ -21,7 +23,7 @@ class NonAutoincrementModel(ormar.Model):
 class ExplicitNullableModel(ormar.Model):
     ormar_config = base_ormar_config.copy()
 
-    id: int = ormar.Integer(primary_key=True, nullable=True)
+    id: Optional[int] = ormar.Integer(primary_key=True, nullable=True)
 
 
 create_test_database = init_tests(base_ormar_config)