Sneak in some utilities

matthiask · matthiask · commit e263efc23500 · 2023-09-08T16:00:34.000+02:00
diff --git a/feincms3_sites/fields.py b/feincms3_sites/fields.py
@@ -0,0 +1,50 @@
+import inspect
+
+from django import forms
+from tree_queries.fields import TreeNodeForeignKey
+
+
+def _variable_from_stack(name, must_exist=()):
+    """
+    We want to filter objects related to a particular site so that content
+    managers can only select objects from the same site as the object which is
+    currently being edited. However, there's no sensible way of passing the
+    edited object down to the form or choice field without doing it by hand for
+    each and every model admin and inline instance. This isn't just tedious:
+    The bigger problem is that errors will happen. Therefore, this helper can
+    be used to fetch a variable from any parent frame. Specifying only the
+    variable name feels a bit brittle, therefore the additional
+    ``must_exist`` argument can be used to require several other variables
+    to exist in the matching scope. This still isn't 100% safe but it works
+    well enough.
+    """
+
+    _sentinel = object()
+    # Start inspecting from the caller of the function which called us
+    frame = inspect.currentframe().f_back.f_back
+    try:
+        while frame:
+            obj = frame.f_locals.get(name, _sentinel)
+            if obj is not _sentinel and all(v in frame.f_locals for v in must_exist):
+                return obj
+            frame = frame.f_back
+    finally:
+        # Delete the frame reference to make the GC's job easier
+        del frame
+
+
+class AutomaticSiteRestrictionChoiceField(forms.ModelChoiceField):
+    def __init__(self, queryset, *args, **kwargs):
+        obj = _variable_from_stack("obj", ("object_id", "to_field"))
+        queryset = queryset.filter(site_id=obj.site_id if obj else None)
+        super().__init__(queryset, *args, **kwargs)
+
+
+class AutomaticSiteRestrictionForeignKey(TreeNodeForeignKey):
+    """
+    The Site foreign key field has to be called ``site`` on all participating objects
+    """
+
+    def formfield(self, **kwargs):
+        kwargs.setdefault("form_class", AutomaticSiteRestrictionChoiceField)
+        return super().formfield(**kwargs)
