add columns to csv for bug budgets, as well as totals submitted/paid

[utils.git] / src / budget_sync / budget_graph.py

diff --git a/src/budget_sync/budget_graph.py b/src/budget_sync/budget_graph.py
index c9319f9f6ea405c51daeef6fdc1a093bdeae45ce..e2824ee30ea4dc70589a610a1cdb74df46692fe0 100644 (file)
--- a/src/budget_sync/budget_graph.py
+++ b/src/budget_sync/budget_graph.py
@@ -16,6 +16,7 @@ from collections import OrderedDict
 
 import collections
 
+
 class OrderedSet(collections.MutableSet):
 
     def __init__(self, iterable=None):
@@ -350,6 +351,22 @@ class Node:
             retval[key] = Payment._from_toml(self, key, value)
         return retval
 
+    @cached_property
+    def submitted_excluding_subtasks(self) -> Money:
+        retval = Money()
+        for payment in self.payments.values():
+            if payment.submitted is not None or payment.paid is not None:
+                retval += payment.amount
+        return retval
+
+    @cached_property
+    def paid_excluding_subtasks(self) -> Money:
+        retval = Money()
+        for payment in self.payments.values():
+            if payment.paid is not None:
+                retval += payment.amount
+        return retval
+
     @property
     def parent(self) -> Optional["Node"]:
         if self.parent_id is not None:
@@ -547,6 +564,7 @@ class BudgetGraphIncorrectRootForMilestone(BudgetGraphError):
 
 class BudgetGraph:
     nodes: Dict[int, Node]
+    milestone_payments: Dict[Milestone, List[Payment]]
 
     def __init__(self, bugs: Iterable[Bug], config: Config):
         self.nodes = OrderedDict()
@@ -559,7 +577,7 @@ class BudgetGraph:
             node.parent.immediate_children.add(node)
         self.milestone_payments = OrderedDict()
         # useful debug prints
-        #for bug in bugs:
+        # for bug in bugs:
         #    node = self.nodes[bug.id]
         #    print ("bug added", bug.id, node, node.parent.immediate_children)
 
@@ -621,7 +639,7 @@ class BudgetGraph:
             subtasks_total += child.fixed_budget_including_subtasks
             childlist.append(child.bug.id)
         # useful debug prints
-        #print ("subtask total", node.bug.id, root.bug.id, subtasks_total,
+        # print ("subtask total", node.bug.id, root.bug.id, subtasks_total,
         #                        childlist)
 
         payees_total = Money(0)
@@ -637,9 +655,9 @@ class BudgetGraph:
                 previous_payment = payee_payments.get(payment.payee)
                 if previous_payment is not None:
                     # NOT AN ERROR
-                    print ("NOT AN ERROR", BudgetGraphDuplicatePayeesForTask(
-                           node.bug.id, root.bug.id,
-                           previous_payment[-1].payee_key, payment.payee_key))
+                    print("NOT AN ERROR", BudgetGraphDuplicatePayeesForTask(
+                        node.bug.id, root.bug.id,
+                        previous_payment[-1].payee_key, payment.payee_key))
                     payee_payments[payment.payee].append(payment)
                 else:
                     payee_payments[payment.payee] = [payment]
@@ -782,6 +800,15 @@ class BudgetGraph:
             retval[node.assignee].append(node)
         return retval
 
+    @cached_property
+    def assigned_nodes_for_milestones(self) -> Dict[Milestone, List[Node]]:
+        retval = {milestone: []
+                  for milestone in self.config.milestones.values()}
+        for node in self.nodes.values():
+            if node.milestone is not None:
+                retval[node.milestone].append(node)
+        return retval
+
     @cached_property
     def payments(self) -> Dict[Person, Dict[Milestone, List[Payment]]]:
         retval = OrderedDict()
@@ -789,7 +816,7 @@ class BudgetGraph:
             milestone_payments = OrderedDict()
             for milestone in self.config.milestones.values():
                 milestone_payments[milestone] = []      # per-person payments
-                self.milestone_payments[milestone] = [] # global payments
+                self.milestone_payments[milestone] = []  # global payments
             retval[person] = milestone_payments
         for node in self.nodes.values():
             if node.milestone is not None:
@@ -799,10 +826,10 @@ class BudgetGraph:
                     self.milestone_payments[node.milestone].append(payment)
         return retval
 
-    def get_milestone_people(self):
+    def get_milestone_people(self) -> Dict[Milestone, OrderedSet]:
         """get a list of people associated with each milestone
         """
-        payments = list(self.payments) # just activate the payments
+        payments = list(self.payments)  # just activate the payments
         retval = OrderedDict()
         for milestone in self.milestone_payments.keys():
             retval[milestone] = OrderedSet()