arthexis 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

core/user_data.py

@@ -5,134 +5,268 @@ from io import BytesIO
 from zipfile import ZipFile
 import json
 
+from django.apps import apps
 from django.conf import settings
-from django.contrib import admin
-from django.contrib import messages
-from django.contrib.contenttypes.fields import GenericForeignKey
-from django.contrib.contenttypes.models import ContentType
-from django.core.exceptions import ValidationError
+from django.contrib import admin, messages
+from django.contrib.auth import get_user_model
+from django.contrib.auth.signals import user_logged_in
 from django.core.management import call_command
-from django.db import IntegrityError, models
-from django.db.models.signals import post_delete, post_save
+from django.db.models.signals import post_save
 from django.dispatch import receiver
+from django.http import HttpResponse, HttpResponseRedirect
 from django.template.response import TemplateResponse
 from django.urls import path, reverse
-from django.http import HttpResponse, HttpResponseRedirect
 from django.utils.translation import gettext as _
 
 from .entity import Entity
 
 
-class UserDatum(models.Model):
-    """Link an :class:`Entity` instance to a user for persistence."""
+def _data_root(user=None) -> Path:
+    path = Path(getattr(user, "data_path", "") or Path(settings.BASE_DIR) / "data")
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
 
-    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)
-    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)
-    object_id = models.PositiveIntegerField()
-    entity = GenericForeignKey("content_type", "object_id")
-    created_on = models.DateTimeField(auto_now_add=True)
+def _username_for(user) -> str:
+    username = ""
+    if hasattr(user, "get_username"):
+        username = user.get_username()
+    if not username and hasattr(user, "username"):
+        username = user.username
+    if not username and getattr(user, "pk", None):
+        username = str(user.pk)
+    return username
 
-    class Meta:
-        unique_together = ("user", "content_type", "object_id")
-        verbose_name = "User Datum"
-        verbose_name_plural = "User Data"
 
+def _user_allows_user_data(user) -> bool:
+    return bool(user) and not getattr(user, "is_profile_restricted", False)
 
-# ---- Fixture utilities ---------------------------------------------------
 
-def _data_dir() -> Path:
-    path = Path(settings.BASE_DIR) / "data"
-    path.mkdir(exist_ok=True)
+def _data_dir(user) -> Path:
+    username = _username_for(user)
+    if not username:
+        raise ValueError("Cannot determine username for fixture directory")
+    path = _data_root(user) / username
+    path.mkdir(parents=True, exist_ok=True)
     return path
 
 
 def _fixture_path(user, instance) -> Path:
-    ct = ContentType.objects.get_for_model(instance)
-    filename = f"{user.pk}_{ct.app_label}_{ct.model}_{instance.pk}.json"
-    return _data_dir() / filename
-
-
-_seed_fixture_cache: dict[tuple[str, int], Path] | None = None
+    model_meta = instance._meta.concrete_model._meta
+    filename = f"{model_meta.app_label}_{model_meta.model_name}_{instance.pk}.json"
+    return _data_dir(user) / filename
 
 
 def _seed_fixture_path(instance) -> Path | None:
-    """Return the fixture file containing the given seed datum."""
-    global _seed_fixture_cache
-    if _seed_fixture_cache is None:
-        _seed_fixture_cache = {}
-        base = Path(settings.BASE_DIR)
-        for path in base.rglob("fixtures/*.json"):
-            try:
-                with path.open() as f:
-                    data = json.load(f)
-            except Exception:
-                continue
-            for obj in data:
-                model = obj.get("model")
-                pk = obj.get("pk")
-                if model and pk is not None:
-                    _seed_fixture_cache[(model, pk)] = path
-    key = (f"{instance._meta.app_label}.{instance._meta.model_name}", instance.pk)
-    return _seed_fixture_cache.get(key)
-
-
-def dump_user_fixture(instance, user) -> None:
-    path = _fixture_path(user, instance)
-    app_label = instance._meta.app_label
-    model_name = instance._meta.model_name
+    label = f"{instance._meta.app_label}.{instance._meta.model_name}"
+    base = Path(settings.BASE_DIR)
+    for path in base.glob("**/fixtures/*.json"):
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+        if not isinstance(data, list) or not data:
+            continue
+        obj = data[0]
+        if obj.get("model") != label:
+            continue
+        pk = obj.get("pk")
+        if pk is not None and pk == instance.pk:
+            return path
+        fields = obj.get("fields", {}) or {}
+        comparable_fields = {
+            key: value
+            for key, value in fields.items()
+            if key not in {"is_seed_data", "is_deleted", "is_user_data"}
+        }
+        if comparable_fields:
+            match = True
+            for field_name, value in comparable_fields.items():
+                if not hasattr(instance, field_name):
+                    match = False
+                    break
+                if getattr(instance, field_name) != value:
+                    match = False
+                    break
+            if match:
+                return path
+    return None
+
+
+def _resolve_fixture_user(instance, fallback=None):
+    UserModel = get_user_model()
+    owner = getattr(instance, "user", None)
+    if isinstance(owner, UserModel):
+        return owner
+    if hasattr(instance, "owner"):
+        try:
+            owner_value = instance.owner
+        except Exception:
+            owner_value = None
+        if isinstance(owner_value, UserModel):
+            return owner_value
+    return fallback
+
+
+def dump_user_fixture(instance, user=None) -> None:
+    model = instance._meta.concrete_model
+    UserModel = get_user_model()
+    if issubclass(UserModel, Entity) and isinstance(instance, UserModel):
+        return
+    target_user = user or _resolve_fixture_user(instance)
+    if target_user is None:
+        return
+    allow_user_data = _user_allows_user_data(target_user)
+    if not allow_user_data:
+        is_user_data = getattr(instance, "is_user_data", False)
+        if not is_user_data and instance.pk:
+            stored_flag = (
+                type(instance)
+                .all_objects.filter(pk=instance.pk)
+                .values_list("is_user_data", flat=True)
+                .first()
+            )
+            is_user_data = bool(stored_flag)
+        if not is_user_data:
+            return
+    meta = model._meta
+    path = _fixture_path(target_user, instance)
     call_command(
         "dumpdata",
-        f"{app_label}.{model_name}",
+        f"{meta.app_label}.{meta.model_name}",
         indent=2,
         pks=str(instance.pk),
         output=str(path),
+        use_natural_foreign_keys=True,
     )
 
 
-def delete_user_fixture(instance, user) -> None:
-    _fixture_path(user, instance).unlink(missing_ok=True)
-
+def delete_user_fixture(instance, user=None) -> None:
+    target_user = user or _resolve_fixture_user(instance)
+    if target_user is None:
+        return
+    _fixture_path(target_user, instance).unlink(missing_ok=True)
 
-# ---- Signals -------------------------------------------------------------
 
-@receiver(post_save)
-def _entity_saved(sender, instance, **kwargs):
-    if isinstance(instance, UserDatum):
+def _mark_fixture_user_data(path: Path) -> None:
+    try:
+        content = path.read_text(encoding="utf-8")
+    except UnicodeDecodeError:
+        try:
+            content = path.read_bytes().decode("latin-1")
+        except Exception:
+            return
+    except Exception:
         return
     try:
-        ct = ContentType.objects.get_for_model(instance)
-        qs = list(UserDatum.objects.filter(content_type=ct, object_id=instance.pk))
+        data = json.loads(content)
     except Exception:
         return
-    for ud in qs:
-        dump_user_fixture(instance, ud.user)
+    if not isinstance(data, list):
+        return
+    for obj in data:
+        label = obj.get("model")
+        if not label:
+            continue
+        try:
+            model = apps.get_model(label)
+        except LookupError:
+            continue
+        if not issubclass(model, Entity):
+            continue
+        pk = obj.get("pk")
+        if pk is None:
+            continue
+        model.all_objects.filter(pk=pk).update(is_user_data=True)
 
 
-@receiver(post_delete)
-def _entity_deleted(sender, instance, **kwargs):
-    if isinstance(instance, UserDatum):
-        return
+def _load_fixture(path: Path, *, mark_user_data: bool = True) -> bool:
+    """Load a fixture from *path* and optionally flag loaded entities."""
+
+    text = None
     try:
-        ct = ContentType.objects.get_for_model(instance)
-        qs = list(UserDatum.objects.filter(content_type=ct, object_id=instance.pk))
+        text = path.read_text(encoding="utf-8")
+    except UnicodeDecodeError:
+        try:
+            text = path.read_bytes().decode("latin-1")
+        except Exception:
+            return False
+        path.write_text(text, encoding="utf-8")
     except Exception:
+        # Continue without cached text so ``call_command`` can surface the
+        # underlying error just as before.
+        pass
+
+    if text is not None:
+        try:
+            data = json.loads(text)
+        except Exception:
+            data = None
+        else:
+            if isinstance(data, list) and not data:
+                path.unlink(missing_ok=True)
+                return False
+
+    try:
+        call_command("loaddata", str(path), ignorenonexistent=True)
+    except Exception:
+        return False
+
+    if mark_user_data:
+        _mark_fixture_user_data(path)
+
+    return True
+
+
+def _fixture_sort_key(path: Path) -> tuple[int, str]:
+    parts = path.name.split("_", 2)
+    model_part = parts[1].lower() if len(parts) >= 2 else ""
+    is_user = model_part == "user"
+    return (0 if is_user else 1, path.name)
+
+
+def _is_user_fixture(path: Path) -> bool:
+    parts = path.name.split("_", 2)
+    return len(parts) >= 2 and parts[1].lower() == "user"
+
+
+_shared_fixtures_loaded = False
+
+
+def load_shared_user_fixtures(*, force: bool = False, user=None) -> None:
+    global _shared_fixtures_loaded
+    if _shared_fixtures_loaded and not force:
         return
-    for ud in qs:
-        delete_user_fixture(instance, ud.user)
-        ud.delete()
+    root = _data_root(user)
+    paths = sorted(root.glob("*.json"), key=_fixture_sort_key)
+    for path in paths:
+        if _is_user_fixture(path):
+            continue
+        _load_fixture(path)
+    _shared_fixtures_loaded = True
 
 
-@receiver(post_save, sender=UserDatum)
-def _userdatum_saved(sender, instance, **kwargs):
-    dump_user_fixture(instance.entity, instance.user)
+def load_user_fixtures(user, *, include_shared: bool = False) -> None:
+    if include_shared:
+        load_shared_user_fixtures(user=user)
+    paths = sorted(_data_dir(user).glob("*.json"), key=_fixture_sort_key)
+    for path in paths:
+        if _is_user_fixture(path):
+            continue
+        _load_fixture(path)
 
 
-@receiver(post_delete, sender=UserDatum)
-def _userdatum_deleted(sender, instance, **kwargs):
-    delete_user_fixture(instance.entity, instance.user)
+@receiver(user_logged_in)
+def _on_login(sender, request, user, **kwargs):
+    load_user_fixtures(user, include_shared=not _shared_fixtures_loaded)
 
 
-# ---- Admin integration ---------------------------------------------------
+@receiver(post_save, sender=settings.AUTH_USER_MODEL)
+def _on_user_created(sender, instance, created, **kwargs):
+    if created:
+        load_shared_user_fixtures(force=True, user=instance)
+        load_user_fixtures(instance)
+
 
 class UserDatumAdminMixin(admin.ModelAdmin):
     """Mixin adding a *User Datum* checkbox to change forms."""
@@ -140,19 +274,27 @@ class UserDatumAdminMixin(admin.ModelAdmin):
     def render_change_form(
         self, request, context, add=False, change=False, form_url="", obj=None
     ):
-        context["show_user_datum"] = True
+        context["show_user_datum"] = issubclass(
+            self.model, Entity
+        ) and _user_allows_user_data(request.user)
+        context["show_seed_datum"] = issubclass(self.model, Entity)
         context["show_save_as_copy"] = issubclass(self.model, Entity) or hasattr(
             self.model, "clone"
         )
         if obj is not None:
-            ct = ContentType.objects.get_for_model(obj)
-            context["is_user_datum"] = UserDatum.objects.filter(
-                user=request.user, content_type=ct, object_id=obj.pk
-            ).exists()
+            context["is_user_datum"] = getattr(obj, "is_user_data", False)
+            context["is_seed_datum"] = getattr(obj, "is_seed_data", False)
         else:
             context["is_user_datum"] = False
+            context["is_seed_datum"] = False
         return super().render_change_form(request, context, add, change, form_url, obj)
 
+
+class EntityModelAdmin(UserDatumAdminMixin, admin.ModelAdmin):
+    """ModelAdmin base class for :class:`Entity` models."""
+
+    change_form_template = "admin/user_datum_change_form.html"
+
     def save_model(self, request, obj, form, change):
         copied = "_saveacopy" in request.POST
         if copied:
@@ -161,57 +303,93 @@ class UserDatumAdminMixin(admin.ModelAdmin):
             form.instance = obj
             try:
                 super().save_model(request, obj, form, False)
-            except IntegrityError:
-                self.message_user(
+            except Exception:
+                messages.error(
                     request,
                     _("Unable to save copy. Adjust unique fields and try again."),
-                    messages.ERROR,
                 )
-                raise ValidationError("save_as_copy")
+                raise
         else:
             super().save_model(request, obj, form, change)
+            if isinstance(obj, Entity):
+                type(obj).all_objects.filter(pk=obj.pk).update(
+                    is_seed_data=obj.is_seed_data, is_user_data=obj.is_user_data
+                )
         if copied:
             return
-        ct = ContentType.objects.get_for_model(obj)
+        target_user = _resolve_fixture_user(obj, request.user)
+        allow_user_data = _user_allows_user_data(target_user)
         if request.POST.get("_user_datum") == "on":
-            UserDatum.objects.get_or_create(
-                user=request.user, content_type=ct, object_id=obj.pk
-            )
-            dump_user_fixture(obj, request.user)
-            path = _fixture_path(request.user, obj)
-            self.message_user(request, f"User datum saved to {path}")
-        else:
-            qs = UserDatum.objects.filter(
-                user=request.user, content_type=ct, object_id=obj.pk
-            )
-            if qs.exists():
-                qs.delete()
-                delete_user_fixture(obj, request.user)
+            if allow_user_data:
+                if not obj.is_user_data:
+                    type(obj).all_objects.filter(pk=obj.pk).update(is_user_data=True)
+                    obj.is_user_data = True
+                dump_user_fixture(obj, target_user)
+                handler = getattr(self, "user_datum_saved", None)
+                if callable(handler):
+                    handler(request, obj)
+                path = _fixture_path(target_user, obj)
+                self.message_user(request, f"User datum saved to {path}")
+            else:
+                if obj.is_user_data:
+                    type(obj).all_objects.filter(pk=obj.pk).update(is_user_data=False)
+                    obj.is_user_data = False
+                    delete_user_fixture(obj, target_user)
+                messages.warning(
+                    request,
+                    _("User data is not available for this account."),
+                )
+        elif obj.is_user_data:
+            type(obj).all_objects.filter(pk=obj.pk).update(is_user_data=False)
+            obj.is_user_data = False
+            delete_user_fixture(obj, target_user)
+            handler = getattr(self, "user_datum_deleted", None)
+            if callable(handler):
+                handler(request, obj)
 
 
 def patch_admin_user_datum() -> None:
-    """Mixin all registered admin classes."""
-    for model, model_admin in list(admin.site._registry.items()):
-        if model is UserDatum:
-            continue
-        if isinstance(model_admin, UserDatumAdminMixin):
-            continue
-        admin.site.unregister(model)
+    """Mixin all registered entity admin classes and future registrations."""
+
+    if getattr(admin.site, "_user_datum_patched", False):
+        return
+
+    def _patched(admin_class):
         template = (
-            getattr(model_admin, "change_form_template", None)
-            or "admin/user_datum_change_form.html"
+            getattr(admin_class, "change_form_template", None)
+            or EntityModelAdmin.change_form_template
         )
-        attrs = {"change_form_template": template}
-        Patched = type(
-            f"Patched{model_admin.__class__.__name__}",
-            (UserDatumAdminMixin, model_admin.__class__),
-            attrs,
+        return type(
+            f"Patched{admin_class.__name__}",
+            (EntityModelAdmin, admin_class),
+            {"change_form_template": template},
         )
-        admin.site.register(model, Patched)
+
+    for model, model_admin in list(admin.site._registry.items()):
+        if issubclass(model, Entity) and not isinstance(model_admin, EntityModelAdmin):
+            admin.site.unregister(model)
+            admin.site.register(model, _patched(model_admin.__class__))
+
+    original_register = admin.site.register
+
+    def register(model_or_iterable, admin_class=None, **options):
+        models = model_or_iterable
+        if not isinstance(models, (list, tuple, set)):
+            models = [models]
+        admin_class = admin_class or admin.ModelAdmin
+        patched_class = admin_class
+        for model in models:
+            if issubclass(model, Entity) and not issubclass(
+                patched_class, EntityModelAdmin
+            ):
+                patched_class = _patched(patched_class)
+        return original_register(model_or_iterable, patched_class, **options)
+
+    admin.site.register = register
+    admin.site._user_datum_patched = True
 
 
 def _seed_data_view(request):
-    """Display all entities marked as seed data."""
     sections = []
     for model, model_admin in admin.site._registry.items():
         if not issubclass(model, Entity):
@@ -225,14 +403,8 @@ def _seed_data_view(request):
                 f"admin:{obj._meta.app_label}_{obj._meta.model_name}_change",
                 args=[obj.pk],
             )
-            fixture_path = _seed_fixture_path(obj)
-            items.append(
-                {
-                    "url": url,
-                    "label": str(obj),
-                    "fixture": fixture_path.name if fixture_path else "",
-                }
-            )
+            fixture = _seed_fixture_path(obj)
+            items.append({"url": url, "label": str(obj), "fixture": fixture})
         sections.append({"opts": model._meta, "items": items})
     context = admin.site.each_context(request)
     context.update({"title": _("Seed Data"), "sections": sections})
@@ -240,38 +412,33 @@ def _seed_data_view(request):
 
 
 def _user_data_view(request):
-    """Display all user datum entities for the current user."""
-    sections = {}
-    qs = UserDatum.objects.filter(user=request.user).select_related("content_type")
-    for ud in qs:
-        model = ud.content_type.model_class()
-        obj = ud.entity
-        url = reverse(
-            f"admin:{obj._meta.app_label}_{obj._meta.model_name}_change",
-            args=[obj.pk],
-        )
-        section = sections.setdefault(model._meta, [])
-        fixture = _fixture_path(request.user, obj)
-        section.append(
-            {"url": url, "label": str(obj), "fixture": fixture.name}
-        )
-    section_list = [{"opts": opts, "items": items} for opts, items in sections.items()]
+    sections = []
+    for model, model_admin in admin.site._registry.items():
+        if not issubclass(model, Entity):
+            continue
+        objs = model.objects.filter(is_user_data=True)
+        if not objs.exists():
+            continue
+        items = []
+        for obj in objs:
+            url = reverse(
+                f"admin:{obj._meta.app_label}_{obj._meta.model_name}_change",
+                args=[obj.pk],
+            )
+            fixture = _fixture_path(request.user, obj)
+            items.append({"url": url, "label": str(obj), "fixture": fixture})
+        sections.append({"opts": model._meta, "items": items})
     context = admin.site.each_context(request)
     context.update(
-        {
-            "title": _("User Data"),
-            "sections": section_list,
-            "import_export": True,
-        }
+        {"title": _("User Data"), "sections": sections, "import_export": True}
     )
     return TemplateResponse(request, "admin/data_list.html", context)
 
 
 def _user_data_export(request):
-    """Return a zip file containing all fixtures for the current user."""
     buffer = BytesIO()
     with ZipFile(buffer, "w") as zf:
-        for path in _data_dir().glob(f"{request.user.pk}_*.json"):
+        for path in _data_dir(request.user).glob("*.json"):
             zf.write(path, arcname=path.name)
     buffer.seek(0)
     response = HttpResponse(buffer.getvalue(), content_type="application/zip")
@@ -282,41 +449,36 @@ def _user_data_export(request):
 
 
 def _user_data_import(request):
-    """Import fixtures from an uploaded zip file."""
     if request.method == "POST" and request.FILES.get("data_zip"):
         with ZipFile(request.FILES["data_zip"]) as zf:
             paths = []
-            data_dir = _data_dir()
+            data_dir = _data_dir(request.user)
             for name in zf.namelist():
                 if not name.endswith(".json"):
                     continue
+                content = zf.read(name)
                 target = data_dir / name
                 with target.open("wb") as f:
-                    f.write(zf.read(name))
+                    f.write(content)
                 paths.append(target)
         if paths:
-            call_command("loaddata", *[str(p) for p in paths])
-            for p in paths:
-                try:
-                    user_id, app_label, model, obj_id = p.stem.split("_", 3)
-                    ct = ContentType.objects.get_by_natural_key(app_label, model)
-                    UserDatum.objects.get_or_create(
-                        user_id=int(user_id), content_type=ct, object_id=int(obj_id)
-                    )
-                except Exception:
-                    continue
+            for path in paths:
+                _load_fixture(path)
     return HttpResponseRedirect(reverse("admin:user_data"))
 
 
 def patch_admin_user_data_views() -> None:
-    """Add custom admin views for seed and user data listings."""
     original_get_urls = admin.site.get_urls
 
     def get_urls():
         urls = original_get_urls()
         custom = [
-            path("seed-data/", admin.site.admin_view(_seed_data_view), name="seed_data"),
-            path("user-data/", admin.site.admin_view(_user_data_view), name="user_data"),
+            path(
+                "seed-data/", admin.site.admin_view(_seed_data_view), name="seed_data"
+            ),
+            path(
+                "user-data/", admin.site.admin_view(_user_data_view), name="user_data"
+            ),
             path(
                 "user-data/export/",
                 admin.site.admin_view(_user_data_export),