add status field to budget_graph.Node

[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 463c3e1b397b8599fba4e96fd6e6dae380868972..322b2404cce91f886c56d85f2534e37eb6015c54 100644 (file)
--- a/src/budget_sync/budget_graph.py
+++ b/src/budget_sync/budget_graph.py
@@ -1,7 +1,6 @@
 from bugzilla.bug import Bug
-from bugzilla import Bugzilla
 from typing import Set, Dict, Iterable, Optional, List, Union, Any
-from budget_sync.util import all_bugs
+from budget_sync.util import BugStatus
 from budget_sync.money import Money
 from budget_sync.config import Config, Person, Milestone
 from functools import cached_property
@@ -212,6 +211,16 @@ class BudgetGraphUnknownMilestone(BudgetGraphParseError):
             f"#{self.bug_id}: unknown milestone: {self.milestone_str!r}"
 
 
+class BudgetGraphUnknownStatus(BudgetGraphParseError):
+    def __init__(self, bug_id: int, status_str: str):
+        super().__init__(bug_id)
+        self.status_str = status_str
+
+    def __str__(self):
+        return f"failed to parse status field of bug " \
+            f"#{self.bug_id}: unknown status: {self.status_str!r}"
+
+
 class Node:
     graph: "BudgetGraph"
     bug: Bug
@@ -236,6 +245,14 @@ class Node:
         if self.milestone_str == "---":
             self.milestone_str = None
 
+    @property
+    def status(self) -> BugStatus:
+        try:
+            return BugStatus.cast(self.bug.status)
+        except ValueError:
+            new_err = BudgetGraphUnknownStatus(self.bug.id, self.bug.status)
+            raise new_err.with_traceback(sys.exc_info()[2])
+
     @cached_property
     def bug_url(self) -> str:
         return f"{self.graph.config.bugzilla_url_stripped}/show_bug.cgi?" \
@@ -330,6 +347,10 @@ class Node:
             milestone = repr(self.milestone)
         except BudgetGraphBaseError:
             milestone = "<unknown milestone>"
+        try:
+            status = repr(self.status)
+        except BudgetGraphBaseError:
+            status = f"<unknown status: {self.bug.status!r}>"
         immediate_children = []
         for i in self.immediate_children:
             immediate_children.append(_NodeSimpleReprWrapper(i))
@@ -347,7 +368,8 @@ class Node:
                 f"milestone_str={self.milestone_str!r}, "
                 f"milestone={milestone}, "
                 f"immediate_children={immediate_children!r}, "
-                f"payments={payments!r}")
+                f"payments={payments!r}, "
+                f"status={status})")
 
 
 class BudgetGraphError(BudgetGraphBaseError):
@@ -497,6 +519,12 @@ class BudgetGraph:
         except BudgetGraphBaseError as e:
             errors.append(e)
 
+        try:
+            # check for status errors
+            node.status
+        except BudgetGraphBaseError as e:
+            errors.append(e)
+
         if node.milestone_str != root.milestone_str:
             errors.append(BudgetGraphMilestoneMismatch(
                 node.bug.id, root.bug.id))
@@ -672,3 +700,13 @@ class BudgetGraph:
                 for payment in node.payments.values():
                     retval[payment.payee][node.milestone].append(payment)
         return retval
+
+    def __repr__(self):
+        nodes = [*self.nodes.values()]
+        try:
+            roots = [_NodeSimpleReprWrapper(i) for i in self.roots]
+            roots.sort()
+            roots_str = repr(roots)
+        except BudgetGraphBaseError:
+            roots_str = "<failed>"
+        return f"BudgetGraph{{nodes={nodes!r}, roots={roots}}}"