[DEV-14235] update tier count

Author: loreleitrimberger

usaspending_api/references/v2/views/filter_tree/psc_filter_tree.py

@@ -251,52 +251,66 @@ def _path_is_valid(self, path: list) -> bool:
                     return False
         return True
 
-    def get_count(self, tiered_keys: list, id) -> int:
+    def get_count(self, tiered_keys: list, id) -> int | None:
         if len(tiered_keys) == 0:
             if id == "Research and Development":
-                return (
-                    PSC.objects.filter(Q(code__startswith="A"))
-                    .annotate(code_prefix=Substr("code", 1, 2))
-                    .values("code_prefix")
-                    .distinct()
-                    .count()
-                )
+                return self.get_research_and_dev_count("A")
             elif id == "Service":
-                return (
-                    PSC.objects.filter(Q(code__regex=r"^[B-Z]$"))
-                    .annotate(code_prefix=Substr("code", 1, 1))
-                    .values("code_prefix")
-                    .distinct()
-                    .count()
-                )
+                return self.get_service_count(r"^[B-Z]", is_regex=True)
             elif id == "Product":
-                return (
-                    PSC.objects.filter(Q(code__regex=r"^\d\d$"))
-                    .annotate(code_prefix=Substr("code", 1, 2))
-                    .values("code_prefix")
-                    .distinct()
-                    .count()
-                )
+                return self.get_product_count(r"^\d\d", is_regex=True)
         else:
             if id.startswith("A"):
                 return self.get_research_and_dev_count(id)
-            elif len(id) == 1:
-                new_length = 2
-            else:
-                new_length = 4
-            all_codes = (
-                PSC.objects.filter(code__startswith=id)
-                .annotate(code_prefix=Substr("code", 1, new_length), code_len=Length("code"))
-                .filter(code_len__gte=new_length)
-            )
-            return all_codes.values("code_prefix").distinct().count()
+            elif re.search(r"^[B-Z]", id):
+                return self.get_service_count(id)
+            elif re.search(r"^\d\d", id):
+                return self.get_product_count(id)
 
-    def get_research_and_dev_count(self, id):
-        new_length = len(id) + 1
-        filters = [Q(code__startswith=id), ~Q(code__endswith=0)]
+    def get_tier_count(self, tier_length, id, is_regex=False, is_r_and_d=False) -> int:
+        if is_regex:
+            filters = [Q(code__regex=id)]
+        else:
+            filters = [Q(code__startswith=id)]
+        if is_r_and_d:
+            filters.append(~Q(code__endswith=0))
         all_codes = (
             PSC.objects.filter(*filters)
-            .annotate(code_prefix=Substr("code", 1, new_length), code_len=Length("code"))
-            .filter(code_len__gte=new_length)
+            .annotate(code_prefix=Substr("code", 1, tier_length), code_len=Length("code"))
+            .filter(code_len__gte=tier_length)
         )
         return all_codes.values("code_prefix").distinct().count()
+
+    def get_research_and_dev_count(self, id) -> int:
+        # Will always need the lowest tier count
+        count = self.get_tier_count(4, id, is_r_and_d=True)
+        if len(id) == 3:  # ex: AA1
+            return count
+
+        count += self.get_tier_count(3, id, is_r_and_d=True)
+        if len(id) == 2:  # ex: AA
+            return count
+
+        count += self.get_tier_count(2, id, is_r_and_d=True)  # id = Research and Development
+        return count
+
+    def get_service_count(self, id, is_regex=False) -> int:
+        count = self.get_tier_count(4, id, is_regex)
+        if len(id) == 2:  # ex: B5
+            return count
+
+        count += self.get_tier_count(2, id, is_regex)
+
+        if len(id) == 1:  # ex: B
+            return count
+
+        count += self.get_tier_count(1, id, is_regex)  # id = Service
+        return count
+
+    def get_product_count(self, id, is_regex=False) -> int:
+        count = self.get_tier_count(4, id, is_regex)
+        if len(id) == 2:  # ex: 10
+            return count
+
+        count += self.get_tier_count(2, id, is_regex)  # id = Product
+        return count