Merge pull request #768 from hrbonz/wip/relative_level

add relative levels to TreeNodeChoiceFieldMixin
django-mptt · Aug 20, 2021 · 913b0ca · 913b0ca
2 parents ecf2811 + f3dc69f
commit 913b0ca
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 1 deletion.
diff --git a/docs/forms.rst b/docs/forms.rst
@@ -63,6 +63,17 @@ options::
    +-- Child 2.1
    +--+-- Child 2.1.1
 
+The starting level can be set so querysets not including the root object can still be displayed in a convenient way. Use the ``start_level`` argument to set the starting point for levels::
+
+   obj = Category.objects.get(pk=1)
+   category = TreeNodeChoiceField(queryset=obj.get_descendants(),
+                                  start_level=obj.level)
+
+...which for this example would result in a select with the following
+options::
+
+   --- Child 1.1.1
+
 .. _`ModelChoiceField`: https://docs.djangoproject.com/en/dev/ref/forms/fields/#django.forms.ModelChoiceField
 
 ``TreeNodeMultipleChoiceField``

diff --git a/mptt/forms.py b/mptt/forms.py
@@ -23,6 +23,7 @@
 class TreeNodeChoiceFieldMixin:
     def __init__(self, queryset, *args, **kwargs):
         self.level_indicator = kwargs.pop("level_indicator", DEFAULT_LEVEL_INDICATOR)
+        self.start_level = kwargs.pop("start_level", 0)
 
         # if a queryset is supplied, enforce ordering
         if hasattr(queryset, "model"):
@@ -31,8 +32,12 @@ def __init__(self, queryset, *args, **kwargs):
 
         super().__init__(queryset, *args, **kwargs)
 
-    def _get_level_indicator(self, obj):
+    def _get_relative_level(self, obj):
         level = getattr(obj, obj._mptt_meta.level_attr)
+        return level - self.start_level
+
+    def _get_level_indicator(self, obj):
+        level = self._get_relative_level(obj)
         return mark_safe(conditional_escape(self.level_indicator) * level)
 
     def label_from_instance(self, obj):

diff --git a/tests/myapp/test_forms.py b/tests/myapp/test_forms.py
@@ -177,3 +177,45 @@ def test_treenodechoicefield_level_indicator(self):
             '<option value="11">+-- Tactical RPG</option>'
             "</select>",
         )
+
+    def test_treenodechoicefield_relative_level(self):
+        top = Genre.objects.get(pk=2)
+        field = TreeNodeChoiceField(queryset=top.get_descendants())
+        self.assertHTMLEqual(
+            field.widget.render("test", None),
+            '<select name="test">'
+            '<option value="" selected>---------</option>'
+            '<option value="3">------ 2D Platformer</option>'
+            '<option value="4">------ 3D Platformer</option>'
+            '<option value="5">------ 4D Platformer</option>'
+            "</select>",
+        )
+
+        field = TreeNodeChoiceField(
+            queryset=top.get_descendants(include_self=True),
+            start_level=top.level,
+        )
+        self.assertHTMLEqual(
+            field.widget.render("test", None),
+            '<select name="test">'
+            '<option value="" selected>---------</option>'
+            '<option value="2"> Platformer</option>'
+            '<option value="3">--- 2D Platformer</option>'
+            '<option value="4">--- 3D Platformer</option>'
+            '<option value="5">--- 4D Platformer</option>'
+            "</select>",
+        )
+
+        field = TreeNodeChoiceField(
+            queryset=top.get_descendants(),
+            start_level=top.level + 1,
+        )
+        self.assertHTMLEqual(
+            field.widget.render("test", None),
+            '<select name="test">'
+            '<option value="" selected>---------</option>'
+            '<option value="3"> 2D Platformer</option>'
+            '<option value="4"> 3D Platformer</option>'
+            '<option value="5"> 4D Platformer</option>'
+            "</select>",
+        )