Tree fix (#6581)

* PartCategoryTree - add "subcategories" field

* Fix rendering of PartCategoryTree

* Implement similar fixes for StockLocationTree

* Bump API version

InvenTree/InvenTree/api_version.py

@@ -1,11 +1,15 @@
 """InvenTree API version information."""
 
 # InvenTree API version
-INVENTREE_API_VERSION = 176
+INVENTREE_API_VERSION = 177
 """Increment this API version number whenever there is a significant change to the API that any clients need to know about."""
 
 INVENTREE_API_TEXT = """
 
+v177 - 2024-02-27 : https://github.com/inventree/InvenTree/pull/6581
+    - Adds "subcategoies" count to PartCategoryTree serializer
+    - Adds "sublocations" count to StockLocationTree serializer
+
 v176 - 2024-02-26 : https://github.com/inventree/InvenTree/pull/6535
     - Adds the field "plugins_install_disabled" to the Server info API endpoint
 

InvenTree/part/api.py

@@ -308,9 +308,17 @@ class CategoryTree(ListAPI):
 
     filter_backends = ORDER_FILTER
 
+    ordering_fields = ['level', 'name', 'subcategories']
+
     # Order by tree level (top levels first) and then name
     ordering = ['level', 'name']
 
+    def get_queryset(self, *args, **kwargs):
+        """Return an annotated queryset for the CategoryTree endpoint."""
+        queryset = super().get_queryset(*args, **kwargs)
+        queryset = part_serializers.CategoryTree.annotate_queryset(queryset)
+        return queryset
+
 
 class CategoryParameterList(ListCreateAPI):
     """API endpoint for accessing a list of PartCategoryParameterTemplate objects.

InvenTree/part/filters.py

@@ -281,6 +281,28 @@ def annotate_category_parts():
     )
 
 
+def annotate_sub_categories():
+    """Construct a queryset annotation which returns the number of subcategories for each provided category."""
+    subquery = part.models.PartCategory.objects.filter(
+        tree_id=OuterRef('tree_id'),
+        lft__gt=OuterRef('lft'),
+        rght__lt=OuterRef('rght'),
+        level__gt=OuterRef('level'),
+    )
+
+    return Coalesce(
+        Subquery(
+            subquery.annotate(
+                total=Func(F('pk'), function='COUNT', output_field=IntegerField())
+            )
+            .values('total')
+            .order_by()
+        ),
+        0,
+        output_field=IntegerField(),
+    )
+
+
 def filter_by_parameter(queryset, template_id: int, value: str, func: str = ''):
     """Filter the given queryset by a given template parameter.
 

InvenTree/part/serializers.py

@@ -123,7 +123,14 @@ class CategoryTree(InvenTree.serializers.InvenTreeModelSerializer):
         """Metaclass defining serializer fields."""
 
         model = PartCategory
-        fields = ['pk', 'name', 'parent', 'icon', 'structural']
+        fields = ['pk', 'name', 'parent', 'icon', 'structural', 'subcategories']
+
+    subcategories = serializers.IntegerField(label=_('Subcategories'), read_only=True)
+
+    @staticmethod
+    def annotate_queryset(queryset):
+        """Annotate the queryset with the number of subcategories."""
+        return queryset.annotate(subcategories=part.filters.annotate_sub_categories())
 
 
 class PartAttachmentSerializer(InvenTree.serializers.InvenTreeAttachmentSerializer):

InvenTree/stock/api.py

@@ -394,9 +394,17 @@ class StockLocationTree(ListAPI):
 
     filter_backends = ORDER_FILTER
 
+    ordering_fields = ['level', 'name', 'sublocations']
+
     # Order by tree level (top levels first) and then name
     ordering = ['level', 'name']
 
+    def get_queryset(self, *args, **kwargs):
+        """Return annotated queryset for the StockLocationTree endpoint."""
+        queryset = super().get_queryset(*args, **kwargs)
+        queryset = StockSerializers.LocationTreeSerializer.annotate_queryset(queryset)
+        return queryset
+
 
 class StockLocationTypeList(ListCreateAPI):
     """API endpoint for a list of StockLocationType objects.

InvenTree/stock/filters.py

@@ -59,3 +59,25 @@ def annotate_child_items():
         0,
         output_field=IntegerField(),
     )
+
+
+def annotate_sub_locations():
+    """Construct a queryset annotation which returns the number of sub-locations below a certain StockLocation node in a StockLocation tree."""
+    subquery = stock.models.StockLocation.objects.filter(
+        tree_id=OuterRef('tree_id'),
+        lft__gt=OuterRef('lft'),
+        rght__lt=OuterRef('rght'),
+        level__gt=OuterRef('level'),
+    )
+
+    return Coalesce(
+        Subquery(
+            subquery.annotate(
+                count=Func(F('pk'), function='COUNT', output_field=IntegerField())
+            )
+            .values('count')
+            .order_by()
+        ),
+        0,
+        output_field=IntegerField(),
+    )

InvenTree/stock/serializers.py

@@ -858,7 +858,14 @@ class LocationTreeSerializer(InvenTree.serializers.InvenTreeModelSerializer):
         """Metaclass options."""
 
         model = StockLocation
-        fields = ['pk', 'name', 'parent', 'icon', 'structural']
+        fields = ['pk', 'name', 'parent', 'icon', 'structural', 'sublocations']
+
+    sublocations = serializers.IntegerField(label=_('Sublocations'), read_only=True)
+
+    @staticmethod
+    def annotate_queryset(queryset):
+        """Annotate the queryset with the number of sublocations."""
+        return queryset.annotate(sublocations=Count('children'))
 
 
 class LocationSerializer(InvenTree.serializers.InvenTreeTagModelSerializer):

src/frontend/src/components/nav/PartCategoryTree.tsx

@@ -8,7 +8,11 @@ import {
   useMantineTheme
 } from '@mantine/core';
 import { ReactTree, ThemeSettings } from '@naisutech/react-tree';
-import { IconSitemap } from '@tabler/icons-react';
+import {
+  IconChevronDown,
+  IconChevronRight,
+  IconSitemap
+} from '@tabler/icons-react';
 import { useQuery } from '@tanstack/react-query';
 import { useMemo } from 'react';
 import { useNavigate } from 'react-router-dom';
@@ -40,7 +44,8 @@ export function PartCategoryTree({
             return {
               id: category.pk,
               label: category.name,
-              parentId: category.parent
+              parentId: category.parent,
+              children: category.subcategories
             };
           })
         )
@@ -67,6 +72,14 @@ export function PartCategoryTree({
     );
   }
 
+  function renderIcon({ node, open }: { node: any; open?: boolean }) {
+    if (node.children == 0) {
+      return undefined;
+    }
+
+    return open ? <IconChevronDown /> : <IconChevronRight />;
+  }
+
   const mantineTheme = useMantineTheme();
 
   const themes: ThemeSettings = useMemo(() => {
@@ -146,6 +159,7 @@ export function PartCategoryTree({
         <ReactTree
           nodes={treeQuery.data ?? []}
           RenderNode={renderNode}
+          RenderIcon={renderIcon}
           defaultSelectedNodes={selectedCategory ? [selectedCategory] : []}
           showEmptyItems={false}
           theme={mantineTheme.colorScheme}

src/frontend/src/components/nav/StockLocationTree.tsx

@@ -1,7 +1,11 @@
 import { t } from '@lingui/macro';
 import { Drawer, Group, LoadingOverlay, Stack, Text } from '@mantine/core';
 import { ReactTree } from '@naisutech/react-tree';
-import { IconSitemap } from '@tabler/icons-react';
+import {
+  IconChevronDown,
+  IconChevronRight,
+  IconSitemap
+} from '@tabler/icons-react';
 import { useQuery } from '@tanstack/react-query';
 import { useNavigate } from 'react-router-dom';
 
@@ -32,7 +36,8 @@ export function StockLocationTree({
             return {
               id: location.pk,
               label: location.name,
-              parentId: location.parent
+              parentId: location.parent,
+              children: location.sublocations
             };
           })
         )
@@ -59,6 +64,14 @@ export function StockLocationTree({
     );
   }
 
+  function renderIcon({ node, open }: { node: any; open?: boolean }) {
+    if (node.children == 0) {
+      return undefined;
+    }
+
+    return open ? <IconChevronDown /> : <IconChevronRight />;
+  }
+
   return (
     <Drawer
       opened={opened}
@@ -87,6 +100,7 @@ export function StockLocationTree({
           nodes={treeQuery.data ?? []}
           showEmptyItems={false}
           RenderNode={renderNode}
+          RenderIcon={renderIcon}
           defaultSelectedNodes={selectedLocation ? [selectedLocation] : []}
         />
       </Stack>