From 85a40ec41877d5f10555cf568395279a4b04c6a8 Mon Sep 17 00:00:00 2001
From: Oliver <oliver.henry.walters@gmail.com>
Date: Wed, 21 Jul 2021 21:23:30 +1000
Subject: [PATCH] Tree exclusion for PartCategory and StockLocation

---
 InvenTree/part/api.py  | 14 ++++++++++++++
 InvenTree/stock/api.py | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/InvenTree/part/api.py b/InvenTree/part/api.py
index 527a9395ee..28e5d4300a 100644
--- a/InvenTree/part/api.py
+++ b/InvenTree/part/api.py
@@ -105,6 +105,20 @@ class CategoryList(generics.ListCreateAPIView):
             except (ValueError, PartCategory.DoesNotExist):
                 pass
 
+        # Exclude PartCategory tree
+        exclude_tree = params.get('exclude_tree', None)
+
+        if exclude_tree is not None:
+            try:
+                cat = PartCategory.objects.get(pk=exclude_tree)
+
+                queryset = queryset.exclude(
+                    pk__in=[c.pk for c in cat.get_descendants(include_self=True)]
+                )
+
+            except (ValueError, PartCategory.DoesNotExist):
+                pass
+
         return queryset
 
     filter_backends = [
diff --git a/InvenTree/stock/api.py b/InvenTree/stock/api.py
index 56df35b5c3..ce5e902cff 100644
--- a/InvenTree/stock/api.py
+++ b/InvenTree/stock/api.py
@@ -343,6 +343,20 @@ class StockLocationList(generics.ListCreateAPIView):
             except (ValueError, StockLocation.DoesNotExist):
                 pass
 
+        # Exclude StockLocation tree
+        exclude_tree = params.get('exclude_tree', None)
+
+        if exclude_tree is not None:
+            try:
+                loc = StockLocation.objects.get(pk=exclude_tree)
+
+                queryset = queryset.exclude(
+                    pk__in=[l.pk for l in loc.get_descendants(include_self=True)]
+                )
+
+            except (ValueError, StockLocation.DoesNotExist):
+                pass
+
         return queryset
 
     filter_backends = [