From 678157aac4d99be38d3a43dd9f884e93c6406234 Mon Sep 17 00:00:00 2001
From: Oliver Walters <oliver.henry.walters@gmail.com>
Date: Sun, 8 Sep 2019 19:13:13 +1000
Subject: [PATCH] Update StockLocation and PartCategory models

- Use the MPTT functionality once more
---
 InvenTree/part/models.py  | 27 +++++++++++++++++++--------
 InvenTree/stock/models.py | 31 +++++++++++++++++++++++--------
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/InvenTree/part/models.py b/InvenTree/part/models.py
index 4590a062a8..2970129d6e 100644
--- a/InvenTree/part/models.py
+++ b/InvenTree/part/models.py
@@ -64,21 +64,32 @@ class PartCategory(InvenTreeTree):
         verbose_name = "Part Category"
         verbose_name_plural = "Part Categories"
 
+    def get_parts(self, cascade=True):
+        """ Return a queryset for all parts under this category.
+
+        args:
+            cascade - If True, also look under subcategories (default = True)
+        """
+
+        if cascade:
+            """ Select any parts which exist in this category or any child categories """
+            query = Part.objects.filter(category__in=self.getUniqueChildren(include_self=True))
+        else:
+            query = Part.objects.filter(category=self.pk)
+
+        return query
+
+
     @property
     def item_count(self):
         return self.partcount()
 
-    def partcount(self, cascade=True, active=True):
+    def partcount(self, cascade=True, active=False):
         """ Return the total part count under this category
         (including children of child categories)
         """
 
-        cats = [self.id]
-
-        if cascade:
-            cats += [cat for cat in self.getUniqueChildren()]
-
-        query = Part.objects.filter(category__in=cats)
+        query = self.get_parts(cascade=cascade)
 
         if active:
             query = query.filter(active=True)
@@ -88,7 +99,7 @@ class PartCategory(InvenTreeTree):
     @property
     def has_parts(self):
         """ True if there are any parts in this category """
-        return self.parts.count() > 0
+        return self.partcount() > 0
 
 
 @receiver(pre_delete, sender=PartCategory, dispatch_uid='partcategory_delete_log')
diff --git a/InvenTree/stock/models.py b/InvenTree/stock/models.py
index 78cef49371..889eaac5bd 100644
--- a/InvenTree/stock/models.py
+++ b/InvenTree/stock/models.py
@@ -34,9 +34,6 @@ class StockLocation(InvenTreeTree):
     def get_absolute_url(self):
         return reverse('stock-location-detail', kwargs={'pk': self.id})
 
-    def has_items(self):
-        return self.stock_items.count() > 0
-
     def format_barcode(self):
         """ Return a JSON string for formatting a barcode for this StockLocation object """
 
@@ -49,16 +46,34 @@ class StockLocation(InvenTreeTree):
             }
         )
 
+    def get_stock_items(self, cascade=True):
+        """ Return a queryset for all stock items under this category.
+
+        Args:
+            cascade: If True, also look under sublocations (default = True)
+        """
+
+        if cascade:
+            query = StockItem.objects.filter(location__in=self.getUniqueChildren(include_self=True))
+        else:
+            query = StockItem.objects.filter(location=self.pk)
+
+        return query
+
     def stock_item_count(self, cascade=True):
         """ Return the number of StockItem objects which live in or under this category
         """
 
-        if cascade:
-            query = StockItem.objects.filter(location__in=self.getUniqueChildren())
-        else:
-            query = StockItem.objects.filter(location=self)
+        return self.get_stock_items(cascade).count()
+
+    def has_items(self, cascade=True):
+        """ Return True if there are StockItems existing in this category.
+
+        Args:
+            cascade: If True, also search an sublocations (default = True)
+        """
+        return self.stock_item_count(cascade) > 0
 
-        return query.count()
 
     @property
     def item_count(self):