fix: session def

ocadotechnology · Nov 7, 2024 · 3b0409e · 3b0409e
1 parent 045deef
commit 3b0409e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 15 deletions.
diff --git a/codeforlife/models/abstract_base_session.py b/codeforlife/models/abstract_base_session.py
@@ -5,7 +5,6 @@
 
 import typing as t
 
-from django.contrib.auth import get_user_model
 from django.contrib.sessions.base_session import (
     AbstractBaseSession as _AbstractBaseSession,
 )
@@ -15,13 +14,17 @@
 
 from .abstract_base_user import AbstractBaseUser
 
+# pylint: disable=duplicate-code
 if t.TYPE_CHECKING:
     from django_stubs_ext.db.models import TypedModelMeta
 
     from .base_session_store import BaseSessionStore
 else:
     TypedModelMeta = object
 
+AnyAbstractBaseUser = t.TypeVar("AnyAbstractBaseUser", bound=AbstractBaseUser)
+# pylint: enable=duplicate-code
+
 
 class AbstractBaseSession(_AbstractBaseSession):
     """
@@ -32,12 +35,6 @@ class AbstractBaseSession(_AbstractBaseSession):
     pk: str  # type: ignore[assignment]
 
     user_id: int
-    user = models.OneToOneField(
-        t.cast(t.Type[AbstractBaseUser], get_user_model()),
-        null=True,
-        blank=True,
-        on_delete=models.CASCADE,
-    )
 
     # pylint: disable-next=missing-class-docstring,too-few-public-methods
     class Meta(TypedModelMeta):
@@ -58,3 +55,24 @@ def store(self):
     @classmethod
     def get_session_store_class(cls) -> t.Type["BaseSessionStore"]:
         raise NotImplementedError
+
+    @staticmethod
+    def init_user_field(user_class: t.Type[AnyAbstractBaseUser]):
+        """Initializes the user field that relates a session to a user.
+
+        Example:
+            class Session(AbstractBaseSession):
+                user = AbstractBaseSession.init_user_field(User)
+
+        Args:
+            user_class: The user model to associate sessions to.
+
+        Returns:
+            A one-to-one field that relates to the provided user model.
+        """
+        return models.OneToOneField(
+            user_class,
+            null=True,
+            blank=True,
+            on_delete=models.CASCADE,
+        )
diff --git a/codeforlife/user/models/session.py b/codeforlife/user/models/session.py
@@ -5,7 +5,6 @@
 
 import typing as t
 
-from django.db import models
 from django.db.models.query import QuerySet
 
 from ...models import AbstractBaseSession, BaseSessionStore
@@ -23,13 +22,7 @@ class Session(AbstractBaseSession):
 
     auth_factors: QuerySet["SessionAuthFactor"]
 
-    # TODO: remove in new schema
-    user = models.OneToOneField(  # type: ignore[assignment]
-        User,
-        null=True,
-        blank=True,
-        on_delete=models.CASCADE,
-    )
+    user = AbstractBaseSession.init_user_field(User)
 
     @classmethod
     def get_session_store_class(cls):