Add type stubs and import for model fields

Author: davegaeddert

example/app/users/models.py

@@ -1,9 +1,10 @@
 from plain import models
+from plain.models import types
 from plain.passwords.models import PasswordField
 
 
 @models.register_model
 class User(models.Model):
-    email = models.EmailField()
+    email: str = types.EmailField()
     password = PasswordField()
-    is_admin = models.BooleanField(default=False)
+    is_admin: bool = types.BooleanField(default=False)

plain-admin/tests/app/users/models.py

@@ -1,7 +1,8 @@
 from plain import models
+from plain.models import types
 
 
 @models.register_model
 class User(models.Model):
-    username = models.CharField(max_length=255)
-    is_admin = models.BooleanField(default=False)
+    username: str = types.CharField(max_length=255)
+    is_admin: bool = types.BooleanField(default=False)

plain-api/plain/api/models.py

@@ -1,8 +1,11 @@
 import binascii
 import os
 import uuid
+from datetime import datetime
+from uuid import UUID
 
 from plain import models
+from plain.models import types
 
 
 def generate_token() -> str:
@@ -11,17 +14,17 @@ def generate_token() -> str:
 
 @models.register_model
 class APIKey(models.Model):
-    uuid = models.UUIDField(default=uuid.uuid4)
-    created_at = models.DateTimeField(auto_now_add=True)
-    updated_at = models.DateTimeField(auto_now=True)
-    expires_at = models.DateTimeField(required=False, allow_null=True)
-    last_used_at = models.DateTimeField(required=False, allow_null=True)
+    uuid: UUID = types.UUIDField(default=uuid.uuid4)
+    created_at: datetime = types.DateTimeField(auto_now_add=True)
+    updated_at: datetime = types.DateTimeField(auto_now=True)
+    expires_at: datetime | None = types.DateTimeField(required=False, allow_null=True)
+    last_used_at: datetime | None = types.DateTimeField(required=False, allow_null=True)
 
-    name = models.CharField(max_length=255, required=False)
+    name: str = types.CharField(max_length=255, required=False)
 
-    token = models.CharField(max_length=40, default=generate_token)
+    token: str = types.CharField(max_length=40, default=generate_token)
 
-    api_version = models.CharField(max_length=255, required=False)
+    api_version: str = types.CharField(max_length=255, required=False)
 
     model_options = models.Options(
         constraints=[

plain-auth/tests/app/users/models.py

@@ -1,7 +1,8 @@
 from plain import models
+from plain.models import types
 
 
 @models.register_model
 class User(models.Model):
-    username = models.CharField(max_length=255)
-    is_admin = models.BooleanField(default=False)
+    username: str = types.CharField(max_length=255)
+    is_admin: bool = types.BooleanField(default=False)

plain-cache/plain/cache/models.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
-from typing import Self
+from datetime import datetime
+from typing import Any, Self
 
 from plain import models
+from plain.models import types
 from plain.utils import timezone
 
 
@@ -19,11 +21,11 @@ def forever(self) -> Self:
 
 @models.register_model
 class CachedItem(models.Model):
-    key = models.CharField(max_length=255)
-    value = models.JSONField(required=False, allow_null=True)
-    expires_at = models.DateTimeField(required=False, allow_null=True)
-    created_at = models.DateTimeField(auto_now_add=True)
-    updated_at = models.DateTimeField(auto_now=True)
+    key: str = types.CharField(max_length=255)
+    value: Any = types.JSONField(required=False, allow_null=True)
+    expires_at: datetime | None = types.DateTimeField(required=False, allow_null=True)
+    created_at: datetime = types.DateTimeField(auto_now_add=True)
+    updated_at: datetime = types.DateTimeField(auto_now=True)
 
     query = CachedItemQuerySet()
 

plain-flags/plain/flags/models.py

@@ -1,7 +1,11 @@
+from __future__ import annotations
+
 import re
+from datetime import datetime
 
 from plain import models
 from plain.exceptions import ValidationError
+from plain.models import types
 
 
 def validate_flag_name(value: str) -> None:
@@ -11,11 +15,11 @@ def validate_flag_name(value: str) -> None:
 
 @models.register_model
 class FlagResult(models.Model):
-    created_at = models.DateTimeField(auto_now_add=True)
-    updated_at = models.DateTimeField(auto_now=True)
-    flag = models.ForeignKey("Flag", on_delete=models.CASCADE)
-    key = models.CharField(max_length=255)
-    value = models.JSONField()
+    created_at: datetime = types.DateTimeField(auto_now_add=True)
+    updated_at: datetime = types.DateTimeField(auto_now=True)
+    flag: Flag = types.ForeignKey("Flag", on_delete=models.CASCADE)
+    key: str = types.CharField(max_length=255)
+    value = types.JSONField()
 
     model_options = models.Options(
         constraints=[
@@ -31,19 +35,19 @@ def __str__(self) -> str:
 
 @models.register_model
 class Flag(models.Model):
-    created_at = models.DateTimeField(auto_now_add=True)
-    updated_at = models.DateTimeField(auto_now=True)
-    name = models.CharField(max_length=255, validators=[validate_flag_name])
+    created_at: datetime = types.DateTimeField(auto_now_add=True)
+    updated_at: datetime = types.DateTimeField(auto_now=True)
+    name: str = types.CharField(max_length=255, validators=[validate_flag_name])
 
     # Optional description that can be filled in after the flag is used/created
-    description = models.TextField(required=False)
+    description: str = types.TextField(required=False)
 
     # To manually disable a flag before completing deleting
     # (good to disable first to make sure the code doesn't use the flag anymore)
-    enabled = models.BooleanField(default=True)
+    enabled: bool = types.BooleanField(default=True)
 
     # To provide an easier way to see if a flag is still being used
-    used_at = models.DateTimeField(required=False, allow_null=True)
+    used_at: datetime | None = types.DateTimeField(required=False, allow_null=True)
 
     model_options = models.Options(
         constraints=[

plain-jobs/plain/jobs/models.py

@@ -3,8 +3,8 @@
 import datetime
 import logging
 import traceback
-import uuid
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING, Any, Self
+from uuid import UUID, uuid4
 
 from opentelemetry import trace
 from opentelemetry.semconv._incubating.attributes.code_attributes import (
@@ -23,7 +23,7 @@
 from opentelemetry.trace import Link, SpanContext, SpanKind
 
 from plain import models
-from plain.models import transaction
+from plain.models import transaction, types
 from plain.models.expressions import F
 from plain.runtime import settings
 from plain.utils import timezone
@@ -44,25 +44,31 @@ class JobRequest(models.Model):
     Keep all pending job requests in a single table.
     """
 
-    created_at = models.DateTimeField(auto_now_add=True)
-    uuid = models.UUIDField(default=uuid.uuid4)
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
+    uuid: UUID = types.UUIDField(default=uuid4)
 
-    job_class = models.CharField(max_length=255)
-    parameters = models.JSONField(required=False, allow_null=True)
-    priority = models.SmallIntegerField(default=0)
-    source = models.TextField(required=False)
-    queue = models.CharField(default="default", max_length=255)
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
 
-    retries = models.SmallIntegerField(default=0)
-    retry_attempt = models.SmallIntegerField(default=0)
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
 
-    concurrency_key = models.CharField(max_length=255, required=False)
+    concurrency_key: str = types.CharField(max_length=255, required=False)
 
-    start_at = models.DateTimeField(required=False, allow_null=True)
+    start_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
 
     # OpenTelemetry trace context
-    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
-    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
 
     # expires_at = models.DateTimeField(required=False, allow_null=True)
 
@@ -149,24 +155,30 @@ class JobProcess(models.Model):
     All active jobs are stored in this table.
     """
 
-    uuid = models.UUIDField(default=uuid.uuid4)
-    created_at = models.DateTimeField(auto_now_add=True)
-    started_at = models.DateTimeField(required=False, allow_null=True)
+    uuid: UUID = types.UUIDField(default=uuid4)
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
+    started_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
 
     # From the JobRequest
-    job_request_uuid = models.UUIDField()
-    job_class = models.CharField(max_length=255)
-    parameters = models.JSONField(required=False, allow_null=True)
-    priority = models.SmallIntegerField(default=0)
-    source = models.TextField(required=False)
-    queue = models.CharField(default="default", max_length=255)
-    retries = models.SmallIntegerField(default=0)
-    retry_attempt = models.SmallIntegerField(default=0)
-    concurrency_key = models.CharField(max_length=255, required=False)
+    job_request_uuid: UUID = types.UUIDField()
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
+    concurrency_key: str = types.CharField(max_length=255, required=False)
 
     # OpenTelemetry trace context
-    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
-    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
 
     query = JobQuerySet()
 
@@ -229,7 +241,7 @@ def run(self) -> JobResult:
             self.save(update_fields=["started_at"])
 
             try:
-                job = jobs_registry.load_job(self.job_class, self.parameters)
+                job = jobs_registry.load_job(self.job_class, self.parameters or {})
                 job.job_process = self
 
                 try:
@@ -446,36 +458,46 @@ class JobResult(models.Model):
     All in-process and completed jobs are stored in this table.
     """
 
-    uuid = models.UUIDField(default=uuid.uuid4)
-    created_at = models.DateTimeField(auto_now_add=True)
+    uuid: UUID = types.UUIDField(default=uuid4)
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
 
     # From the Job
-    job_process_uuid = models.UUIDField()
-    started_at = models.DateTimeField(required=False, allow_null=True)
-    ended_at = models.DateTimeField(required=False, allow_null=True)
-    error = models.TextField(required=False)
-    status = models.CharField(
+    job_process_uuid: UUID = types.UUIDField()
+    started_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+    ended_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+    error: str = types.TextField(required=False)
+    status: str = types.CharField(
         max_length=20,
         choices=JobResultStatuses.choices,
     )
 
     # From the JobRequest
-    job_request_uuid = models.UUIDField()
-    job_class = models.CharField(max_length=255)
-    parameters = models.JSONField(required=False, allow_null=True)
-    priority = models.SmallIntegerField(default=0)
-    source = models.TextField(required=False)
-    queue = models.CharField(default="default", max_length=255)
-    retries = models.SmallIntegerField(default=0)
-    retry_attempt = models.SmallIntegerField(default=0)
-    concurrency_key = models.CharField(max_length=255, required=False)
+    job_request_uuid: UUID = types.UUIDField()
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
+    concurrency_key: str = types.CharField(max_length=255, required=False)
 
     # Retries
-    retry_job_request_uuid = models.UUIDField(required=False, allow_null=True)
+    retry_job_request_uuid: UUID | None = types.UUIDField(
+        required=False, allow_null=True
+    )
 
     # OpenTelemetry trace context
-    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
-    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
 
     query = JobResultQuerySet()
 
@@ -502,7 +524,7 @@ class JobResult(models.Model):
 
     def retry_job(self, delay: int | None = None) -> JobRequest | None:
         retry_attempt = self.retry_attempt + 1
-        job = jobs_registry.load_job(self.job_class, self.parameters)
+        job = jobs_registry.load_job(self.job_class, self.parameters or {})
 
         if delay is None:
             retry_delay = job.calculate_retry_delay(retry_attempt)

plain-models/plain/models/__init__.py

@@ -41,6 +41,7 @@
 from .options import Options
 from .query import QuerySet
 from .query_utils import Q
+from . import types
 
 # This module exports the user-facing API for defining model classes,
 # with a secondary focus on the most common query utilities like Q.
@@ -104,4 +105,6 @@
     # From registry
     "register_model",
     "models_registry",
+    # Typed field imports
+    "types",
 ]

plain-models/plain/models/base.py

@@ -4,7 +4,7 @@
 import warnings
 from collections.abc import Iterable, Iterator, Sequence
 from itertools import chain
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, dataclass_transform
 
 if TYPE_CHECKING:
     from plain.models.meta import Meta
@@ -48,6 +48,7 @@ def __str__(self) -> str:
 DEFERRED = Deferred()
 
 
+@dataclass_transform(kw_only_default=True)
 class ModelBase(type):
     """Metaclass for all models."""