src/budget_sync/budget_graph.py

   1 from budget_sync.ordered_set import OrderedSet
   2 from bugzilla.bug import Bug
   3 from typing import Callable, Set, Dict, Iterable, Optional, List, Tuple, Union, Any
   4 from budget_sync.util import BugStatus, PrettyPrinter
   5 from budget_sync.money import Money
   6 from budget_sync.config import Config, Person, Milestone
   7 import toml
   8 import sys
   9 import enum
  10 from collections import deque
  11 from datetime import date, time, datetime
  12 try:
  13     from functools import cached_property
  14 except ImportError:
  15     # compatability with python < 3.8
  16     from cached_property import cached_property
  17
  18
  19 class BudgetGraphBaseError(Exception):
  20     pass
  21
  22
  23 class BudgetGraphParseError(BudgetGraphBaseError):
  24     def __init__(self, bug_id: int):
  25         self.bug_id = bug_id
  26
  27
  28 class BudgetGraphPayeesParseError(BudgetGraphParseError):
  29     def __init__(self, bug_id: int, msg: str):
  30         super().__init__(bug_id)
  31         self.msg = msg
  32
  33     def __str__(self):
  34         return f"Failed to parse cf_payees_list field of " \
  35             f"bug #{self.bug_id}: {self.msg}"
  36
  37
  38 class BudgetGraphUnknownAssignee(BudgetGraphParseError):
  39     def __init__(self, bug_id: int, assignee: str):
  40         super().__init__(bug_id)
  41         self.assignee = assignee
  42
  43     def __str__(self):
  44         return f"Bug #{self.bug_id} is assigned to an unknown person: " \
  45             f"{self.assignee!r}"
  46
  47
  48 class BudgetGraphLoopError(BudgetGraphBaseError):
  49     def __init__(self, bug_ids: List[int]):
  50         self.bug_ids = bug_ids
  51
  52     def __str__(self):
  53         retval = f"Detected Loop in Budget Graph: #{self.bug_ids[-1]} -> "
  54         retval += " -> ".join((f"#{i}" for i in self.bug_ids))
  55         return retval
  56
  57
  58 class _NodeSimpleReprWrapper:
  59     def __init__(self, node: "Node"):
  60         self.node = node
  61
  62     def __repr__(self):
  63         return f"#{self.node.bug.id}"
  64
  65     def __lt__(self, other):
  66         # for list.sort()
  67         return self.node.bug.id < other.node.bug.id
  68
  69
  70 class PayeeState(enum.Enum):
  71     NotYetSubmitted = "not yet submitted"
  72     Submitted = "submitted"
  73     Paid = "paid"
  74
  75
  76 _Date = Union[date, datetime]
  77
  78
  79 def _parse_money_from_toml(value: Any) -> Money:
  80     if not isinstance(value, (int, str)):
  81         msg = f"monetary amount is not a string or integer " \
  82             f"(to use fractional amounts such as 123.45, write " \
  83             f"\"123.45\"): {value!r}"
  84         raise ValueError(msg)
  85     return Money(value)
  86
  87
  88 def _parse_date_time_or_none_from_toml(value: Any) -> Optional[_Date]:
  89     if value is None or isinstance(value, (date, datetime)):
  90         return value
  91     elif isinstance(value, time):
  92         msg = f"just a time of day by itself is not enough," \
  93             f" a date must be included: {str(value)}"
  94         raise ValueError(msg)
  95     elif isinstance(value, bool):
  96         msg = f"invalid date: {str(value).lower()}"
  97         raise ValueError(msg)
  98     elif isinstance(value, (str, int, float)):
  99         msg = f"invalid date: {value!r}"
 100         raise ValueError(msg)
 101     else:
 102         msg = f"invalid date"
 103         raise ValueError(msg)
 104
 105
 106 class Payment:
 107     def __init__(self,
 108                  node: "Node",
 109                  payee_key: str,
 110                  amount: Money,
 111                  paid: Optional[_Date],
 112                  submitted: Optional[_Date]):
 113         self.node = node
 114         self.payee_key = payee_key
 115         self.amount = amount
 116         self.paid = paid
 117         self.submitted = submitted
 118
 119     @cached_property
 120     def payee(self) -> Person:
 121         try:
 122             return self.node.graph.config.all_names[self.payee_key]
 123         except KeyError:
 124             msg = f"unknown payee name: {self.payee_key!r} is not the name " \
 125                 f"or an alias of any known person"
 126             raise BudgetGraphPayeesParseError(self.node.bug.id, msg) \
 127                 .with_traceback(sys.exc_info()[2])
 128
 129     @property
 130     def state(self):
 131         if self.paid is not None:
 132             return PayeeState.Paid
 133         if self.submitted is not None:
 134             return PayeeState.Submitted
 135         return PayeeState.NotYetSubmitted
 136
 137     @staticmethod
 138     def _from_toml(node: "Node", payee_key: str, toml_value: Any) -> "Payment":
 139         paid = None
 140         submitted = None
 141         known_keys = ("paid", "submitted", "amount")
 142         if isinstance(toml_value, dict):
 143             try:
 144                 amount = toml_value['amount']
 145             except KeyError:
 146                 msg = f"value for key {payee_key!r} is missing the " \
 147                     f"`amount` field which is required"
 148                 raise BudgetGraphPayeesParseError(node.bug.id, msg) \
 149                     .with_traceback(sys.exc_info()[2])
 150             for k, v in toml_value.items():
 151                 if k in ("paid", "submitted"):
 152                     try:
 153                         parsed_value = _parse_date_time_or_none_from_toml(v)
 154                     except ValueError as e:
 155                         msg = f"failed to parse `{k}` field for" \
 156                             f" key {payee_key!r}: {e}"
 157                         raise BudgetGraphPayeesParseError(
 158                             node.bug.id, msg) \
 159                             .with_traceback(sys.exc_info()[2])
 160                     if k == "paid":
 161                         paid = parsed_value
 162                     else:
 163                         assert k == "submitted"
 164                         submitted = parsed_value
 165                 if k not in known_keys:
 166                     msg = f"value for key {payee_key!r} has an unknown" \
 167                         f" field: `{k}`"
 168                     raise BudgetGraphPayeesParseError(node.bug.id, msg) \
 169                         .with_traceback(sys.exc_info()[2])
 170             try:
 171                 paid = _parse_date_time_or_none_from_toml(
 172                     toml_value.get('paid'))
 173             except ValueError as e:
 174                 msg = f"failed to parse `paid` field for" \
 175                     f" key {payee_key!r}: {e}"
 176                 raise BudgetGraphPayeesParseError(
 177                     node.bug.id, msg) \
 178                     .with_traceback(sys.exc_info()[2])
 179             try:
 180                 submitted = _parse_date_time_or_none_from_toml(
 181                     toml_value.get('submitted'))
 182             except ValueError as e:
 183                 msg = f"failed to parse `submitted` field for" \
 184                     f" key {payee_key!r}: {e}"
 185                 raise BudgetGraphPayeesParseError(
 186                     node.bug.id, msg) \
 187                     .with_traceback(sys.exc_info()[2])
 188         elif isinstance(toml_value, (int, str, float)):
 189             # float included for better error messages
 190             amount = toml_value
 191         else:
 192             msg = f"value for key {payee_key!r} is invalid -- it should " \
 193                 f"either be a monetary value or a table"
 194             raise BudgetGraphPayeesParseError(node.bug.id, msg)
 195         try:
 196             amount = _parse_money_from_toml(amount)
 197         except ValueError as e:
 198             msg = f"failed to parse monetary amount for key {payee_key!r}: {e}"
 199             raise BudgetGraphPayeesParseError(
 200                 node.bug.id, msg) \
 201                 .with_traceback(sys.exc_info()[2])
 202         return Payment(node=node, payee_key=payee_key, amount=amount,
 203                        paid=paid, submitted=submitted)
 204
 205     def __repr__(self):
 206         try:
 207             payee = f"Person<{self.payee.identifier!r}>"
 208         except BudgetGraphBaseError:
 209             payee = "<unknown person>"
 210         return (f"Payment(node={_NodeSimpleReprWrapper(self.node)}, "
 211                 f"payee={payee}, "
 212                 f"payee_key={self.payee_key!r}, "
 213                 f"amount={self.amount}, "
 214                 f"state={self.state.name}, "
 215                 f"paid={str(self.paid)}, "
 216                 f"submitted={str(self.submitted)})")
 217
 218
 219 @enum.unique
 220 class PaymentSummaryState(enum.Enum):
 221     Submitted = PayeeState.Submitted
 222     Paid = PayeeState.Paid
 223     NotYetSubmitted = PayeeState.NotYetSubmitted
 224     Inconsistent = None
 225
 226
 227 class PaymentSummary:
 228     total_submitted: Money
 229     """includes amount paid"""
 230
 231     def __init__(self, payments: Iterable[Payment]):
 232         self.payments = tuple(payments)
 233         self.total = Money(0)
 234         self.total_paid = Money(0)
 235         self.total_submitted = Money(0)
 236         self.submitted_date = None
 237         self.paid_date = None
 238         summary_state = None
 239         for payment in self.payments:
 240             if summary_state is None:
 241                 summary_state = PaymentSummaryState(payment.state)
 242                 self.submitted_date = payment.submitted
 243                 self.paid_date = payment.paid
 244             elif summary_state != PaymentSummaryState(payment.state) \
 245                     or self.submitted_date != payment.submitted \
 246                     or self.paid_date != payment.paid:
 247                 summary_state = PaymentSummaryState.Inconsistent
 248                 self.paid_date = None
 249                 self.submitted_date = None
 250             self.total += payment.amount
 251             if payment.state is PayeeState.Submitted:
 252                 self.total_submitted += payment.amount
 253             elif payment.state is PayeeState.Paid:
 254                 self.total_submitted += payment.amount
 255                 self.total_paid += payment.amount
 256             else:
 257                 assert payment.state is PayeeState.NotYetSubmitted
 258         if summary_state is None:
 259             self.state = PaymentSummaryState.NotYetSubmitted
 260         else:
 261             self.state = summary_state
 262
 263     def __repr__(self) -> str:
 264         return (f"PaymentSummary(total={self.total}, "
 265                 f"total_paid={self.total_paid}, "
 266                 f"total_submitted={self.total_submitted}, "
 267                 f"submitted_date={self.submitted_date}, "
 268                 f"paid_date={self.paid_date}, "
 269                 f"state={self.state}, "
 270                 f"payments={self.payments})")
 271
 272     def __pretty_print__(self, pp: PrettyPrinter):
 273         with pp.type_pp("PaymentSummary") as tpp:
 274             tpp.field("total", self.total)
 275             tpp.field("total_submitted", self.total_submitted)
 276             tpp.field("submitted_date", self.submitted_date)
 277             tpp.field("paid_date", self.paid_date)
 278             tpp.field("state", self.state)
 279             tpp.field("payments", self.payments)
 280
 281
 282 class BudgetGraphUnknownMilestone(BudgetGraphParseError):
 283     def __init__(self, bug_id: int, milestone_str: str):
 284         super().__init__(bug_id)
 285         self.milestone_str = milestone_str
 286
 287     def __str__(self):
 288         return f"failed to parse cf_nlnet_milestone field of bug " \
 289             f"#{self.bug_id}: unknown milestone: {self.milestone_str!r}"
 290
 291
 292 class BudgetGraphUnknownStatus(BudgetGraphParseError):
 293     def __init__(self, bug_id: int, status_str: str):
 294         super().__init__(bug_id)
 295         self.status_str = status_str
 296
 297     def __str__(self):
 298         return f"failed to parse status field of bug " \
 299             f"#{self.bug_id}: unknown status: {self.status_str!r}"
 300
 301
 302 class Node:
 303     graph: "BudgetGraph"
 304     bug: Bug
 305     parent_id: Optional[int]
 306     immediate_children: OrderedSet["Node"]
 307     budget_excluding_subtasks: Money
 308     budget_including_subtasks: Money
 309     fixed_budget_excluding_subtasks: Money
 310     fixed_budget_including_subtasks: Money
 311     milestone_str: Optional[str]
 312
 313     def __init__(self, graph: "BudgetGraph", bug: Bug):
 314         self.graph = graph
 315         self.bug = bug
 316         self.parent_id = getattr(bug, "cf_budget_parent", None)
 317         self.immediate_children = OrderedSet()
 318         self.budget_excluding_subtasks = Money.from_str(bug.cf_budget)
 319         self.fixed_budget_excluding_subtasks = self.budget_excluding_subtasks
 320         self.budget_including_subtasks = Money.from_str(bug.cf_total_budget)
 321         self.fixed_budget_including_subtasks = self.budget_including_subtasks
 322         self.milestone_str = bug.cf_nlnet_milestone
 323         if self.milestone_str == "---":
 324             self.milestone_str = None
 325
 326     @property
 327     def status(self) -> BugStatus:
 328         try:
 329             return BugStatus.cast(self.bug.status)
 330         except ValueError:
 331             new_err = BudgetGraphUnknownStatus(self.bug.id, self.bug.status)
 332             raise new_err.with_traceback(sys.exc_info()[2])
 333
 334     @cached_property
 335     def assignee(self) -> Person:
 336         try:
 337             return self.graph.config.all_names[self.bug.assigned_to]
 338         except KeyError:
 339             raise BudgetGraphUnknownAssignee(self.bug.id,
 340                                              self.bug.assigned_to) \
 341                 .with_traceback(sys.exc_info()[2])
 342
 343     @cached_property
 344     def bug_url(self) -> str:
 345         return f"{self.graph.config.bugzilla_url_stripped}/show_bug.cgi?" \
 346             f"id={self.bug.id}"
 347
 348     @cached_property
 349     def milestone(self) -> Optional[Milestone]:
 350         if self.milestone_str is None:
 351             return None
 352         try:
 353             return self.graph.config.milestones[self.milestone_str]
 354         except KeyError:
 355             new_err = BudgetGraphUnknownMilestone(
 356                 self.bug.id, self.milestone_str)
 357             raise new_err.with_traceback(sys.exc_info()[2])
 358
 359     @cached_property
 360     def payments(self) -> Dict[str, Payment]:
 361         try:
 362             parsed = toml.loads(self.bug.cf_payees_list)
 363         except toml.TomlDecodeError as e:
 364             new_err = BudgetGraphPayeesParseError(
 365                 self.bug.id, f"TOML parse error: {e}")
 366             raise new_err.with_traceback(sys.exc_info()[2])
 367         retval = {}
 368         for key, value in parsed.items():
 369             if not isinstance(key, str):
 370                 raise BudgetGraphPayeesParseError(
 371                     self.bug.id, f"key is not a string: {key!r}")
 372             retval[key] = Payment._from_toml(self, key, value)
 373         return retval
 374
 375     @cached_property
 376     def resolved_payments(self) -> Dict[Person, List[Payment]]:
 377         retval: Dict[Person, List[Payment]] = {}
 378         for payment in self.payments.values():
 379             if payment.payee not in retval:
 380                 retval[payment.payee] = []
 381             retval[payment.payee].append(payment)
 382         return retval
 383
 384     @cached_property
 385     def payment_summaries(self) -> Dict[Person, PaymentSummary]:
 386         return {person: PaymentSummary(payments)
 387                 for person, payments in self.resolved_payments.items()}
 388
 389     @cached_property
 390     def submitted_excluding_subtasks(self) -> Money:
 391         retval = Money()
 392         for payment in self.payments.values():
 393             if payment.submitted is not None or payment.paid is not None:
 394                 retval += payment.amount
 395         return retval
 396
 397     @cached_property
 398     def paid_excluding_subtasks(self) -> Money:
 399         retval = Money()
 400         for payment in self.payments.values():
 401             if payment.paid is not None:
 402                 retval += payment.amount
 403         return retval
 404
 405     @property
 406     def parent(self) -> Optional["Node"]:
 407         if self.parent_id is not None:
 408             return self.graph.nodes[self.parent_id]
 409         return None
 410
 411     def parents(self) -> Iterable["Node"]:
 412         parent = self.parent
 413         while parent is not None:
 414             yield parent
 415             parent = parent.parent
 416
 417     def _raise_loop_error(self):
 418         bug_ids = []
 419         for parent in self.parents():
 420             bug_ids.append(parent.bug.id)
 421             if parent == self:
 422                 break
 423         raise BudgetGraphLoopError(bug_ids)
 424
 425     @cached_property
 426     def root(self) -> "Node":
 427         # also checks for loop errors
 428         retval = self
 429         for parent in self.parents():
 430             retval = parent
 431             if parent == self:
 432                 self._raise_loop_error()
 433         return retval
 434
 435     def children(self) -> Iterable["Node"]:
 436         def visitor(node: Node) -> Iterable[Node]:
 437             for i in node.immediate_children:
 438                 yield i
 439                 yield from visitor(i)
 440         return visitor(self)
 441
 442     def children_breadth_first(self) -> Iterable["Node"]:
 443         q = deque(self.immediate_children)
 444         while True:
 445             try:
 446                 node = q.popleft()
 447             except IndexError:
 448                 return
 449             q.extend(node.immediate_children)
 450             yield node
 451
 452     def __eq__(self, other):
 453         return self.bug.id == other.bug.id
 454
 455     def __hash__(self):
 456         return self.bug.id
 457
 458     def __pretty_print__(self, pp: PrettyPrinter):
 459         with pp.type_pp("Node") as tpp:
 460             tpp.field("graph", ...)
 461             tpp.field("id", _NodeSimpleReprWrapper(self))
 462             tpp.try_field("root",
 463                           lambda: _NodeSimpleReprWrapper(self.root),
 464                           BudgetGraphLoopError)
 465             parent = f"#{self.parent_id}" if self.parent_id is not None else None
 466             tpp.field("parent", parent)
 467             tpp.field("budget_excluding_subtasks",
 468                       self.budget_excluding_subtasks)
 469             tpp.field("budget_including_subtasks",
 470                       self.budget_including_subtasks)
 471             tpp.field("fixed_budget_excluding_subtasks",
 472                       self.fixed_budget_excluding_subtasks)
 473             tpp.field("fixed_budget_including_subtasks",
 474                       self.fixed_budget_including_subtasks)
 475             tpp.field("milestone_str", self.milestone_str)
 476             tpp.try_field("milestone", lambda: self.milestone,
 477                           BudgetGraphBaseError)
 478             immediate_children = [_NodeSimpleReprWrapper(i)
 479                                   for i in self.immediate_children]
 480             tpp.field("immediate_children", immediate_children)
 481             tpp.try_field("payments",
 482                           lambda: list(self.payments.values()),
 483                           BudgetGraphBaseError)
 484             try:
 485                 status = repr(self.status)
 486             except BudgetGraphBaseError:
 487                 status = f"<unknown status: {self.bug.status!r}>"
 488             tpp.field("status", status)
 489             try:
 490                 assignee = f"Person<{self.assignee.identifier!r}>"
 491             except BudgetGraphBaseError:
 492                 assignee = f"<unknown assignee: {self.bug.assigned_to!r}>"
 493             tpp.field("assignee", assignee)
 494             tpp.try_field("resolved_payments",
 495                           lambda: self.resolved_payments,
 496                           BudgetGraphBaseError)
 497             tpp.try_field("payment_summaries",
 498                           lambda: self.payment_summaries,
 499                           BudgetGraphBaseError)
 500
 501     def __repr__(self):
 502         try:
 503             root = _NodeSimpleReprWrapper(self.root)
 504         except BudgetGraphLoopError:
 505             root = "<loop error>"
 506         try:
 507             milestone = repr(self.milestone)
 508         except BudgetGraphBaseError:
 509             milestone = "<unknown milestone>"
 510         try:
 511             status = repr(self.status)
 512         except BudgetGraphBaseError:
 513             status = f"<unknown status: {self.bug.status!r}>"
 514         try:
 515             assignee = f"Person<{self.assignee.identifier!r}>"
 516         except BudgetGraphBaseError:
 517             assignee = f"<unknown assignee: {self.bug.assigned_to!r}>"
 518         immediate_children = []
 519         for i in self.immediate_children:
 520             immediate_children.append(_NodeSimpleReprWrapper(i))
 521         immediate_children.sort()
 522         parent = f"#{self.parent_id}" if self.parent_id is not None else None
 523         payments = list(self.payments.values())
 524         resolved_payments = self.resolved_payments
 525         payment_summaries = self.payment_summaries
 526         return (f"Node(graph=..., "
 527                 f"id={_NodeSimpleReprWrapper(self)}, "
 528                 f"root={root}, "
 529                 f"parent={parent}, "
 530                 f"budget_excluding_subtasks={self.budget_excluding_subtasks}, "
 531                 f"budget_including_subtasks={self.budget_including_subtasks}, "
 532                 f"fixed_budget_excluding_subtasks={self.fixed_budget_excluding_subtasks}, "
 533                 f"fixed_budget_including_subtasks={self.fixed_budget_including_subtasks}, "
 534                 f"milestone_str={self.milestone_str!r}, "
 535                 f"milestone={milestone}, "
 536                 f"immediate_children={immediate_children!r}, "
 537                 f"payments={payments!r}, "
 538                 f"status={status}, "
 539                 f"assignee={assignee}, "
 540                 f"resolved_payments={resolved_payments!r}, "
 541                 f"payment_summaries={payment_summaries!r})")
 542
 543
 544 class BudgetGraphError(BudgetGraphBaseError):
 545     def __init__(self, bug_id: int, root_bug_id: int):
 546         self.bug_id = bug_id
 547         self.root_bug_id = root_bug_id
 548
 549
 550 class BudgetGraphMoneyWithNoMilestone(BudgetGraphError):
 551     def __str__(self):
 552         return (f"Bug assigned money but without"
 553                 f" any assigned milestone: #{self.bug_id}")
 554
 555
 556 class BudgetGraphMilestoneMismatch(BudgetGraphError):
 557     def __str__(self):
 558         return (f"Bug's assigned milestone doesn't match the milestone "
 559                 f"assigned to the root bug: descendant bug"
 560                 f" #{self.bug_id}, root bug"
 561                 f" #{self.root_bug_id}")
 562
 563
 564 class BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(BudgetGraphError):
 565     def __init__(self, bug_id: int, root_bug_id: int,
 566                  expected_budget_excluding_subtasks: Money):
 567         super().__init__(bug_id, root_bug_id)
 568         self.expected_budget_excluding_subtasks = \
 569             expected_budget_excluding_subtasks
 570
 571     def __str__(self):
 572         return (f"Budget assigned to task excluding subtasks "
 573                 f"(cf_budget field) doesn't match calculated value: "
 574                 f"bug #{self.bug_id}, calculated value"
 575                 f" {self.expected_budget_excluding_subtasks}")
 576
 577
 578 class BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(BudgetGraphError):
 579     def __init__(self, bug_id: int, root_bug_id: int,
 580                  expected_budget_including_subtasks: Money):
 581         super().__init__(bug_id, root_bug_id)
 582         self.expected_budget_including_subtasks = \
 583             expected_budget_including_subtasks
 584
 585     def __str__(self):
 586         return (f"Budget assigned to task including subtasks "
 587                 f"(cf_total_budget field) doesn't match calculated value: "
 588                 f"bug #{self.bug_id}, calculated value"
 589                 f" {self.expected_budget_including_subtasks}")
 590
 591
 592 class BudgetGraphNegativeMoney(BudgetGraphError):
 593     def __str__(self):
 594         return (f"Budget assigned to task is less than zero: "
 595                 f"bug #{self.bug_id}")
 596
 597
 598 class BudgetGraphPayeesMoneyMismatch(BudgetGraphError):
 599     def __init__(self, bug_id: int, root_bug_id: int, payees_total: Money,
 600                  expected_payees_total: Money):
 601         super().__init__(bug_id, root_bug_id)
 602         self.payees_total = payees_total
 603         self.expected_payees_total = expected_payees_total
 604
 605     def __str__(self):
 606         return (f"Total budget assigned to payees (cf_payees_list) doesn't "
 607                 f"match expected value: bug #{self.bug_id}, calculated total "
 608                 f"{self.payees_total}, expected value "
 609                 f"{self.expected_payees_total}")
 610
 611
 612 class BudgetGraphNegativePayeeMoney(BudgetGraphError):
 613     def __init__(self, bug_id: int, root_bug_id: int, payee_key: str):
 614         super().__init__(bug_id, root_bug_id)
 615         self.payee_key = payee_key
 616
 617     def __str__(self):
 618         return (f"Budget assigned to payee for task is less than zero: "
 619                 f"bug #{self.bug_id}, payee {self.payee_key!r}")
 620
 621
 622 class BudgetGraphIncorrectRootForMilestone(BudgetGraphError):
 623     def __init__(self, bug_id: int, milestone: str, milestone_canonical_bug_id: int):
 624         super().__init__(bug_id, bug_id)
 625         self.milestone = milestone
 626         self.milestone_canonical_bug_id = milestone_canonical_bug_id
 627
 628     def __str__(self):
 629         return (f"Bug #{self.bug_id} is not the canonical root bug for "
 630                 f"assigned milestone {self.milestone!r} but has no parent "
 631                 f"bug set: the milestone's canonical root bug is "
 632                 f"#{self.milestone_canonical_bug_id}")
 633
 634
 635 class BudgetGraph:
 636     nodes: Dict[int, Node]
 637
 638     def __init__(self, bugs: Iterable[Bug], config: Config):
 639         self.nodes = {}
 640         self.config = config
 641         for bug in bugs:
 642             self.nodes[bug.id] = Node(self, bug)
 643         for node in self.nodes.values():
 644             if node.parent is None:
 645                 continue
 646             node.parent.immediate_children.add(node)
 647         # useful debug prints
 648         # for bug in bugs:
 649         #    node = self.nodes[bug.id]
 650         #    print ("bug added", bug.id, node, node.parent.immediate_children)
 651
 652     @cached_property
 653     def roots(self) -> OrderedSet[Node]:
 654         roots = OrderedSet()
 655         for node in self.nodes.values():
 656             # calling .root also checks for loop errors
 657             root = node.root
 658             roots.add(root)
 659         return roots
 660
 661     def _get_node_errors(self, root: Node, node: Node,
 662                          errors: List[BudgetGraphBaseError]):
 663         if node.milestone_str is None:
 664             if node.budget_including_subtasks != 0 \
 665                     or node.budget_excluding_subtasks != 0:
 666                 errors.append(BudgetGraphMoneyWithNoMilestone(
 667                     node.bug.id, root.bug.id))
 668
 669         try:
 670             # check for milestone errors
 671             node.milestone
 672             if root == node and node.milestone is not None \
 673                     and node.milestone.canonical_bug_id != node.bug.id:
 674                 if node.budget_including_subtasks != 0 \
 675                         or node.budget_excluding_subtasks != 0:
 676                     errors.append(BudgetGraphIncorrectRootForMilestone(
 677                         node.bug.id, node.milestone.identifier,
 678                         node.milestone.canonical_bug_id
 679                     ))
 680         except BudgetGraphBaseError as e:
 681             errors.append(e)
 682
 683         try:
 684             # check for status errors
 685             node.status
 686         except BudgetGraphBaseError as e:
 687             errors.append(e)
 688
 689         try:
 690             # check for assignee errors
 691             node.assignee
 692         except BudgetGraphBaseError as e:
 693             errors.append(e)
 694
 695         if node.milestone_str != root.milestone_str:
 696             errors.append(BudgetGraphMilestoneMismatch(
 697                 node.bug.id, root.bug.id))
 698
 699         if node.budget_excluding_subtasks < 0 \
 700                 or node.budget_including_subtasks < 0:
 701             errors.append(BudgetGraphNegativeMoney(
 702                 node.bug.id, root.bug.id))
 703
 704         childlist = []
 705         subtasks_total = Money(0)
 706         for child in node.immediate_children:
 707             subtasks_total += child.fixed_budget_including_subtasks
 708             childlist.append(child.bug.id)
 709         # useful debug prints
 710         # print ("subtask total", node.bug.id, root.bug.id, subtasks_total,
 711         #                        childlist)
 712
 713         payees_total = Money(0)
 714         payee_payments: Dict[Person, List[Payment]] = {}
 715         for payment in node.payments.values():
 716             if payment.amount < 0:
 717                 errors.append(BudgetGraphNegativePayeeMoney(
 718                     node.bug.id, root.bug.id, payment.payee_key))
 719             payees_total += payment.amount
 720             try:
 721                 # check for payee errors
 722                 payment.payee
 723                 previous_payment = payee_payments.get(payment.payee)
 724                 if previous_payment is not None:
 725                     payee_payments[payment.payee].append(payment)
 726                 else:
 727                     payee_payments[payment.payee] = [payment]
 728             except BudgetGraphBaseError as e:
 729                 errors.append(e)
 730
 731         def set_including_from_excluding_and_error():
 732             node.fixed_budget_including_subtasks = \
 733                 node.budget_excluding_subtasks + subtasks_total
 734             errors.append(
 735                 BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(
 736                     node.bug.id, root.bug.id,
 737                     node.fixed_budget_including_subtasks))
 738
 739         def set_including_from_payees_and_error():
 740             node.fixed_budget_including_subtasks = \
 741                 payees_total + subtasks_total
 742             errors.append(
 743                 BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(
 744                     node.bug.id, root.bug.id,
 745                     node.fixed_budget_including_subtasks))
 746
 747         def set_excluding_from_including_and_error():
 748             node.fixed_budget_excluding_subtasks = \
 749                 node.budget_including_subtasks - subtasks_total
 750             errors.append(
 751                 BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(
 752                     node.bug.id, root.bug.id,
 753                     node.fixed_budget_excluding_subtasks))
 754
 755         def set_excluding_from_payees_and_error():
 756             node.fixed_budget_excluding_subtasks = \
 757                 payees_total
 758             errors.append(
 759                 BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(
 760                     node.bug.id, root.bug.id,
 761                     node.fixed_budget_excluding_subtasks))
 762
 763         def set_payees_from_including_and_error():
 764             fixed_payees_total = \
 765                 node.budget_including_subtasks - subtasks_total
 766             errors.append(BudgetGraphPayeesMoneyMismatch(
 767                 node.bug.id, root.bug.id, payees_total, fixed_payees_total))
 768
 769         def set_payees_from_excluding_and_error():
 770             fixed_payees_total = \
 771                 node.budget_excluding_subtasks
 772             errors.append(BudgetGraphPayeesMoneyMismatch(
 773                 node.bug.id, root.bug.id, payees_total, fixed_payees_total))
 774
 775         payees_matches_including = \
 776             node.budget_including_subtasks - subtasks_total == payees_total
 777         payees_matches_excluding = \
 778             node.budget_excluding_subtasks == payees_total
 779         including_matches_excluding = \
 780             node.budget_including_subtasks - subtasks_total \
 781             == node.budget_excluding_subtasks
 782
 783         if payees_matches_including \
 784                 and payees_matches_excluding \
 785                 and including_matches_excluding:
 786             pass  # no error
 787         elif payees_matches_including:
 788             # can't have 2 match without all 3 matching
 789             assert not payees_matches_excluding
 790             assert not including_matches_excluding
 791             if node.budget_including_subtasks == 0 and len(node.payments) == 0:
 792                 set_including_from_excluding_and_error()
 793             else:
 794                 set_excluding_from_including_and_error()
 795         elif payees_matches_excluding:
 796             # can't have 2 match without all 3 matching
 797             assert not payees_matches_including
 798             assert not including_matches_excluding
 799             if node.budget_excluding_subtasks == 0 and len(node.payments) == 0:
 800                 if node.budget_including_subtasks == 0:
 801                     set_including_from_excluding_and_error()
 802                 else:
 803                     set_excluding_from_including_and_error()
 804             else:
 805                 set_including_from_excluding_and_error()
 806         elif including_matches_excluding:
 807             # can't have 2 match without all 3 matching
 808             assert not payees_matches_including
 809             assert not payees_matches_excluding
 810             if len(node.payments) == 0:
 811                 pass  # no error -- payees is just not set
 812             elif node.budget_excluding_subtasks == 0 \
 813                     and node.budget_including_subtasks == 0:
 814                 set_excluding_from_payees_and_error()
 815                 set_including_from_payees_and_error()
 816             else:
 817                 set_payees_from_excluding_and_error()
 818         else:
 819             # nothing matches
 820             if len(node.payments) == 0:
 821                 # payees unset -- don't need to set payees
 822                 if node.budget_including_subtasks == 0:
 823                     set_including_from_excluding_and_error()
 824                 else:
 825                     set_excluding_from_including_and_error()
 826             elif node.budget_excluding_subtasks == 0 \
 827                     and node.budget_including_subtasks == 0:
 828                 set_excluding_from_payees_and_error()
 829                 set_including_from_payees_and_error()
 830             elif node.budget_excluding_subtasks == 0:
 831                 set_excluding_from_including_and_error()
 832                 set_payees_from_including_and_error()
 833             elif node.budget_including_subtasks == 0:
 834                 set_including_from_excluding_and_error()
 835                 set_payees_from_excluding_and_error()
 836             else:
 837                 set_including_from_excluding_and_error()
 838                 set_payees_from_excluding_and_error()
 839
 840     def get_errors(self) -> List[BudgetGraphBaseError]:
 841         errors = []
 842         try:
 843             roots = self.roots
 844         except BudgetGraphBaseError as e:
 845             errors.append(e)
 846             return errors
 847
 848         for root in roots:
 849             try:
 850                 for child in reversed(list(root.children_breadth_first())):
 851                     try:
 852                         self._get_node_errors(root, child, errors)
 853                     except BudgetGraphBaseError as e:
 854                         errors.append(e)
 855                 self._get_node_errors(root, root, errors)
 856             except BudgetGraphBaseError as e:
 857                 errors.append(e)
 858         return errors
 859
 860     @cached_property
 861     def assigned_nodes(self) -> Dict[Person, List[Node]]:
 862         retval: Dict[Person, List[Node]]
 863         retval = {person: [] for person in self.config.people.values()}
 864         for node in self.nodes.values():
 865             retval[node.assignee].append(node)
 866         return retval
 867
 868     @cached_property
 869     def assigned_nodes_for_milestones(self) -> Dict[Milestone, List[Node]]:
 870         retval: Dict[Milestone, List[Node]]
 871         retval = {milestone: []
 872                   for milestone in self.config.milestones.values()}
 873         for node in self.nodes.values():
 874             if node.milestone is not None:
 875                 retval[node.milestone].append(node)
 876         return retval
 877
 878     @cached_property
 879     def milestone_payments(self) -> Dict[Milestone, List[Payment]]:
 880         retval: Dict[Milestone, List[Payment]] = {
 881             milestone: [] for milestone in self.config.milestones.values()
 882         }
 883         for node in self.nodes.values():
 884             if node.milestone is not None:
 885                 retval[node.milestone].extend(node.payments.values())
 886         return retval
 887
 888     @cached_property
 889     def payments(self) -> Dict[Person, Dict[Milestone, List[Payment]]]:
 890         retval: Dict[Person, Dict[Milestone, List[Payment]]] = {
 891             person: {
 892                 milestone: []
 893                 for milestone in self.config.milestones.values()
 894             }
 895             for person in self.config.people.values()
 896         }
 897         for node in self.nodes.values():
 898             if node.milestone is not None:
 899                 for payment in node.payments.values():
 900                     retval[payment.payee][node.milestone].append(payment)
 901         return retval
 902
 903     @cached_property
 904     def milestone_people(self) -> Dict[Milestone, OrderedSet[Person]]:
 905         """get a list of people associated with each milestone
 906         """
 907         payments = list(self.payments)  # just activate the payments
 908         retval = {}
 909         for milestone in self.milestone_payments.keys():
 910             retval[milestone] = OrderedSet()
 911         for milestone, payments in self.milestone_payments.items():
 912             for payment in payments:
 913                 retval[milestone].add(payment.payee)
 914         return retval
 915
 916     def __pretty_print__(self, pp: PrettyPrinter):
 917         with pp.type_pp("BudgetGraph") as tpp:
 918             tpp.field("nodes", self.nodes)
 919             tpp.try_field("roots",
 920                           lambda: [_NodeSimpleReprWrapper(i)
 921                                    for i in self.roots],
 922                           BudgetGraphBaseError)
 923             tpp.try_field("assigned_nodes",
 924                           lambda: {
 925                               person: [
 926                                   _NodeSimpleReprWrapper(node)
 927                                   for node in nodes
 928                               ]
 929                               for person, nodes in self.assigned_nodes.items()
 930                           },
 931                           BudgetGraphBaseError)
 932             tpp.try_field("assigned_nodes_for_milestones",
 933                           lambda: {
 934                               milestone: [
 935                                   _NodeSimpleReprWrapper(node)
 936                                   for node in nodes
 937                               ]
 938                               for milestone, nodes in self.assigned_nodes_for_milestones.items()
 939                           },
 940                           BudgetGraphBaseError)
 941             tpp.try_field("payments",
 942                           lambda: self.payments, BudgetGraphBaseError)
 943             tpp.try_field("milestone_people",
 944                           lambda: self.milestone_people,
 945                           BudgetGraphBaseError)
 946
 947     def __repr__(self):
 948         nodes = [*self.nodes.values()]
 949
 950         def repr_or_failed(f: Callable[[], Any]) -> str:
 951             try:
 952                 return repr(f())
 953             except BudgetGraphBaseError:
 954                 return "<failed>"
 955
 956         try:
 957             roots = [_NodeSimpleReprWrapper(i) for i in self.roots]
 958             roots.sort()
 959             roots_str = repr(roots)
 960         except BudgetGraphBaseError:
 961             roots_str = "<failed>"
 962         assigned_nodes = repr_or_failed(lambda: {
 963             person: [
 964                 _NodeSimpleReprWrapper(node)
 965                 for node in nodes
 966             ]
 967             for person, nodes in self.assigned_nodes.items()
 968         })
 969         assigned_nodes_for_milestones = repr_or_failed(lambda: {
 970             milestone: [
 971                 _NodeSimpleReprWrapper(node)
 972                 for node in nodes
 973             ]
 974             for milestone, nodes in self.assigned_nodes_for_milestones.items()
 975         })
 976         milestone_payments = repr_or_failed(lambda: self.milestone_payments)
 977         payments = repr_or_failed(lambda: self.payments)
 978         milestone_people = repr_or_failed(lambda: self.milestone_people)
 979         return (f"BudgetGraph{{nodes={nodes!r}, "
 980                 f"roots={roots}, "
 981                 f"assigned_nodes={assigned_nodes}, "
 982                 f"assigned_nodes_for_milestones={assigned_nodes_for_milestones}, "
 983                 f"milestone_payments={milestone_payments}, "
 984                 f"payments={payments}, "
 985                 f"milestone_people={milestone_people}}}")