plain.models 0.49.2__py3-none-any.whl → 0.51.0__py3-none-any.whl

plain/models/otel.py

@@ -1,10 +1,18 @@
+from __future__ import annotations
+
 import re
 import traceback
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from opentelemetry import context as otel_context
 from opentelemetry import trace
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+    from plain.models.backends.base.base import BaseDatabaseWrapper
 from opentelemetry.semconv._incubating.attributes.db_attributes import (
     DB_QUERY_PARAMETER_TEMPLATE,
     DB_USER,
@@ -99,7 +107,9 @@ def _clean_identifier(identifier: str) -> str:
 
 
 @contextmanager
-def db_span(db, sql: Any, *, many: bool = False, params=None):
+def db_span(
+    db: BaseDatabaseWrapper, sql: Any, *, many: bool = False, params: Any = None
+) -> Generator[Span | None, None, None]:
     """Open an OpenTelemetry CLIENT span for a database query.
 
     All common attributes (`db.*`, `network.*`, etc.) are set automatically.
@@ -107,7 +117,7 @@ def db_span(db, sql: Any, *, many: bool = False, params=None):
     """
 
     # Fast-exit if instrumentation suppression flag set in context.
-    if otel_context.get_value(_SUPPRESS_KEY):
+    if otel_context.get_value(_SUPPRESS_KEY):  # type: ignore[arg-type]
         yield None
         return
 
@@ -177,15 +187,15 @@ def db_span(db, sql: Any, *, many: bool = False, params=None):
 
 
 @contextmanager
-def suppress_db_tracing():
-    token = otel_context.attach(otel_context.set_value(_SUPPRESS_KEY, True))
+def suppress_db_tracing() -> Generator[None, None, None]:
+    token = otel_context.attach(otel_context.set_value(_SUPPRESS_KEY, True))  # type: ignore[arg-type]
     try:
         yield
     finally:
         otel_context.detach(token)
 
 
-def _get_code_attributes():
+def _get_code_attributes() -> dict[str, Any]:
     """Extract code context attributes for the current database query.
 
     Returns a dict of OpenTelemetry code attributes.

plain/models/preflight.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 import inspect
 from collections import defaultdict
+from collections.abc import Callable
+from typing import Any
 
 from plain.models.db import db_connection
-from plain.models.registry import models_registry
+from plain.models.registry import ModelsRegistry, models_registry
 from plain.packages import packages_registry
 from plain.preflight import PreflightCheck, PreflightResult, register_check
 
@@ -11,7 +15,7 @@ from plain.preflight import PreflightCheck, PreflightResult, register_check
 class CheckDatabaseBackends(PreflightCheck):
     """Validates database backend configuration when plain.models is available."""
 
-    def run(self):
+    def run(self) -> list[PreflightResult]:
         return db_connection.validation.preflight()
 
 
@@ -19,14 +23,16 @@ class CheckDatabaseBackends(PreflightCheck):
 class CheckAllModels(PreflightCheck):
     """Validates all model definitions for common issues."""
 
-    def run(self):
+    def run(self) -> list[PreflightResult]:
         db_table_models = defaultdict(list)
         indexes = defaultdict(list)
         constraints = defaultdict(list)
         errors = []
         models = models_registry.get_models()
         for model in models:
-            db_table_models[model._meta.db_table].append(model._meta.label)
+            db_table_models[model.model_options.db_table].append(
+                model.model_options.label
+            )
             if not inspect.ismethod(model.preflight):
                 errors.append(
                     PreflightResult(
@@ -37,10 +43,10 @@ class CheckAllModels(PreflightCheck):
                 )
             else:
                 errors.extend(model.preflight())
-            for model_index in model._meta.indexes:
-                indexes[model_index.name].append(model._meta.label)
-            for model_constraint in model._meta.constraints:
-                constraints[model_constraint.name].append(model._meta.label)
+            for model_index in model.model_options.indexes:
+                indexes[model_index.name].append(model.model_options.label)
+            for model_constraint in model.model_options.constraints:
+                constraints[model_constraint.name].append(model.model_options.label)
         for db_table, model_labels in db_table_models.items():
             if len(model_labels) != 1:
                 model_labels_str = ", ".join(model_labels)
@@ -84,7 +90,9 @@ class CheckAllModels(PreflightCheck):
         return errors
 
 
-def _check_lazy_references(models_registry, packages_registry):
+def _check_lazy_references(
+    models_registry: ModelsRegistry, packages_registry: Any
+) -> list[PreflightResult]:
     """
     Ensure all lazy (i.e. string) model references have been resolved.
 
@@ -98,7 +106,9 @@ def _check_lazy_references(models_registry, packages_registry):
     if not pending_models:
         return []
 
-    def extract_operation(obj):
+    def extract_operation(
+        obj: Any,
+    ) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
         """
         Take a callable found in Packages._pending_operations and identify the
         original callable passed to Packages.lazy_model_operation(). If that
@@ -115,7 +125,7 @@ def _check_lazy_references(models_registry, packages_registry):
             operation = operation.func
         return operation, args, keywords
 
-    def app_model_error(model_key):
+    def app_model_error(model_key: tuple[str, str]) -> str:
         try:
             packages_registry.get_package_config(model_key[0])
             model_error = "app '{}' doesn't provide model '{}'".format(*model_key)
@@ -129,7 +139,12 @@ def _check_lazy_references(models_registry, packages_registry):
     # pair, the original lazy function, and its positional and keyword args as
     # determined by extract_operation().
 
-    def field_error(model_key, func, args, keywords):
+    def field_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult:
         error_msg = (
             "The field %(field)s was declared with a lazy reference "
             "to '%(model)s', but %(model_error)s."
@@ -145,7 +160,12 @@ def _check_lazy_references(models_registry, packages_registry):
             id="fields.lazy_reference_not_resolvable",
         )
 
-    def default_error(model_key, func, args, keywords):
+    def default_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult:
         error_msg = (
             "%(op)s contains a lazy reference to %(model)s, but %(model_error)s."
         )
@@ -167,8 +187,13 @@ def _check_lazy_references(models_registry, packages_registry):
         ("plain.models.fields.related", "resolve_related_class"): field_error,
     }
 
-    def build_error(model_key, func, args, keywords):
-        key = (func.__module__, func.__name__)
+    def build_error(
+        model_key: tuple[str, str],
+        func: Callable[..., Any],
+        args: list[Any],
+        keywords: dict[str, Any],
+    ) -> PreflightResult | None:
+        key = (func.__module__, func.__name__)  # type: ignore[attr-defined]
         error_fn = known_lazy.get(key, default_error)
         return error_fn(model_key, func, args, keywords) if error_fn else None
 
@@ -189,7 +214,7 @@ def _check_lazy_references(models_registry, packages_registry):
 class CheckLazyReferences(PreflightCheck):
     """Ensures all lazy (string) model references have been resolved."""
 
-    def run(self):
+    def run(self) -> list[PreflightResult]:
         return _check_lazy_references(models_registry, packages_registry)
 
 
@@ -197,7 +222,7 @@ class CheckLazyReferences(PreflightCheck):
 class CheckDatabaseTables(PreflightCheck):
     """Checks for unknown tables in the database when plain.models is available."""
 
-    def run(self):
+    def run(self) -> list[PreflightResult]:
         errors = []
 
         db_tables = db_connection.introspection.table_names()