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         self.not_submitted = []
 239         summary_state = None
 240         for payment in self.payments:
 241             if summary_state is None:
 242                 summary_state = PaymentSummaryState(payment.state)
 243                 self.submitted_date = payment.submitted
 244                 self.paid_date = payment.paid
 245             elif summary_state != PaymentSummaryState(payment.state) \
 246                     or self.submitted_date != payment.submitted \
 247                     or self.paid_date != payment.paid:
 248                 summary_state = PaymentSummaryState.Inconsistent
 249                 self.paid_date = None
 250                 self.submitted_date = None
 251             self.total += payment.amount
 252             if payment.state is PayeeState.Submitted:
 253                 self.total_submitted += payment.amount
 254             elif payment.state is PayeeState.Paid:
 255                 self.total_submitted += payment.amount
 256                 self.total_paid += payment.amount
 257             else:
 258                 assert payment.state is PayeeState.NotYetSubmitted
 259                 self.not_submitted.append(payment.node.bug.id)
 260         if summary_state is None:
 261             self.state = PaymentSummaryState.NotYetSubmitted
 262         else:
 263             self.state = summary_state
 264
 265     def get_not_submitted(self):
 266         return self.not_submitted
 267
 268     def __repr__(self) -> str:
 269         return (f"PaymentSummary(total={self.total}, "
 270                 f"total_paid={self.total_paid}, "
 271                 f"total_submitted={self.total_submitted}, "
 272                 f"submitted_date={self.submitted_date}, "
 273                 f"paid_date={self.paid_date}, "
 274                 f"state={self.state}, "
 275                 f"payments={self.payments})")
 276
 277     def __pretty_print__(self, pp: PrettyPrinter):
 278         with pp.type_pp("PaymentSummary") as tpp:
 279             tpp.field("total", self.total)
 280             tpp.field("total_submitted", self.total_submitted)
 281             tpp.field("submitted_date", self.submitted_date)
 282             tpp.field("paid_date", self.paid_date)
 283             tpp.field("state", self.state)
 284             tpp.field("payments", self.payments)
 285
 286
 287 class BudgetGraphUnknownMilestone(BudgetGraphParseError):
 288     def __init__(self, bug_id: int, milestone_str: str):
 289         super().__init__(bug_id)
 290         self.milestone_str = milestone_str
 291
 292     def __str__(self):
 293         return f"failed to parse cf_nlnet_milestone field of bug " \
 294             f"#{self.bug_id}: unknown milestone: {self.milestone_str!r}"
 295
 296
 297 class BudgetGraphUnknownStatus(BudgetGraphParseError):
 298     def __init__(self, bug_id: int, status_str: str):
 299         super().__init__(bug_id)
 300         self.status_str = status_str
 301
 302     def __str__(self):
 303         return f"failed to parse status field of bug " \
 304             f"#{self.bug_id}: unknown status: {self.status_str!r}"
 305
 306
 307 class Node:
 308     graph: "BudgetGraph"
 309     bug: Bug
 310     parent_id: Optional[int]
 311     immediate_children: OrderedSet["Node"]
 312     budget_excluding_subtasks: Money
 313     budget_including_subtasks: Money
 314     fixed_budget_excluding_subtasks: Money
 315     fixed_budget_including_subtasks: Money
 316     milestone_str: Optional[str]
 317
 318     def __init__(self, graph: "BudgetGraph", bug: Bug):
 319         self.graph = graph
 320         self.bug = bug
 321         self.parent_id = getattr(bug, "cf_budget_parent", None)
 322         self.immediate_children = OrderedSet()
 323         self.budget_excluding_subtasks = Money.from_str(bug.cf_budget)
 324         self.fixed_budget_excluding_subtasks = self.budget_excluding_subtasks
 325         self.budget_including_subtasks = Money.from_str(bug.cf_total_budget)
 326         self.fixed_budget_including_subtasks = self.budget_including_subtasks
 327         self.milestone_str = bug.cf_nlnet_milestone
 328         if self.milestone_str == "---":
 329             self.milestone_str = None
 330
 331     @property
 332     def status(self) -> BugStatus:
 333         try:
 334             return BugStatus.cast(self.bug.status)
 335         except ValueError:
 336             new_err = BudgetGraphUnknownStatus(self.bug.id, self.bug.status)
 337             raise new_err.with_traceback(sys.exc_info()[2])
 338
 339     @cached_property
 340     def assignee(self) -> Person:
 341         try:
 342             return self.graph.config.all_names[self.bug.assigned_to]
 343         except KeyError:
 344             raise BudgetGraphUnknownAssignee(self.bug.id,
 345                                              self.bug.assigned_to) \
 346                 .with_traceback(sys.exc_info()[2])
 347
 348     @cached_property
 349     def bug_url(self) -> str:
 350         return f"{self.graph.config.bugzilla_url_stripped}/show_bug.cgi?" \
 351             f"id={self.bug.id}"
 352
 353     @cached_property
 354     def milestone(self) -> Optional[Milestone]:
 355         if self.milestone_str is None:
 356             return None
 357         try:
 358             return self.graph.config.milestones[self.milestone_str]
 359         except KeyError:
 360             new_err = BudgetGraphUnknownMilestone(
 361                 self.bug.id, self.milestone_str)
 362             raise new_err.with_traceback(sys.exc_info()[2])
 363
 364     @cached_property
 365     def payments(self) -> Dict[str, Payment]:
 366         try:
 367             parsed = toml.loads(self.bug.cf_payees_list)
 368         except toml.TomlDecodeError as e:
 369             new_err = BudgetGraphPayeesParseError(
 370                 self.bug.id, f"TOML parse error: {e}")
 371             raise new_err.with_traceback(sys.exc_info()[2])
 372         retval = {}
 373         for key, value in parsed.items():
 374             if not isinstance(key, str):
 375                 raise BudgetGraphPayeesParseError(
 376                     self.bug.id, f"key is not a string: {key!r}")
 377             retval[key] = Payment._from_toml(self, key, value)
 378         return retval
 379
 380     @cached_property
 381     def resolved_payments(self) -> Dict[Person, List[Payment]]:
 382         retval: Dict[Person, List[Payment]] = {}
 383         for payment in self.payments.values():
 384             if payment.payee not in retval:
 385                 retval[payment.payee] = []
 386             retval[payment.payee].append(payment)
 387         return retval
 388
 389     @cached_property
 390     def payment_summaries(self) -> Dict[Person, PaymentSummary]:
 391         return {person: PaymentSummary(payments)
 392                 for person, payments in self.resolved_payments.items()}
 393
 394     @cached_property
 395     def submitted_excluding_subtasks(self) -> Money:
 396         retval = Money()
 397         for payment in self.payments.values():
 398             if payment.submitted is not None or payment.paid is not None:
 399                 retval += payment.amount
 400         return retval
 401
 402     @cached_property
 403     def paid_excluding_subtasks(self) -> Money:
 404         retval = Money()
 405         for payment in self.payments.values():
 406             if payment.paid is not None:
 407                 retval += payment.amount
 408         return retval
 409
 410     @cached_property
 411     def submitted_including_subtasks(self) -> Money:
 412         retval = self.submitted_excluding_subtasks
 413         for i in self.immediate_children:
 414             retval += i.submitted_including_subtasks
 415         return retval
 416
 417     @cached_property
 418     def paid_including_subtasks(self) -> Money:
 419         retval = self.paid_excluding_subtasks
 420         for i in self.immediate_children:
 421             retval += i.paid_including_subtasks
 422         return retval
 423
 424     @property
 425     def parent(self) -> Optional["Node"]:
 426         if self.parent_id is not None:
 427             return self.graph.nodes[self.parent_id]
 428         return None
 429
 430     def parents(self) -> Iterable["Node"]:
 431         parent = self.parent
 432         while parent is not None:
 433             yield parent
 434             parent = parent.parent
 435
 436     def _raise_loop_error(self):
 437         bug_ids = []
 438         for parent in self.parents():
 439             bug_ids.append(parent.bug.id)
 440             if parent == self:
 441                 break
 442         raise BudgetGraphLoopError(bug_ids)
 443
 444     @cached_property
 445     def root(self) -> "Node":
 446         # also checks for loop errors
 447         retval = self
 448         for parent in self.parents():
 449             retval = parent
 450             if parent == self:
 451                 self._raise_loop_error()
 452         return retval
 453
 454     def children(self) -> Iterable["Node"]:
 455         def visitor(node: Node) -> Iterable[Node]:
 456             for i in node.immediate_children:
 457                 yield i
 458                 yield from visitor(i)
 459         return visitor(self)
 460
 461     def children_breadth_first(self) -> Iterable["Node"]:
 462         q = deque(self.immediate_children)
 463         while True:
 464             try:
 465                 node = q.popleft()
 466             except IndexError:
 467                 return
 468             q.extend(node.immediate_children)
 469             yield node
 470
 471     def __eq__(self, other):
 472         return self.bug.id == other.bug.id
 473
 474     def __hash__(self):
 475         return self.bug.id
 476
 477     def __pretty_print__(self, pp: PrettyPrinter):
 478         with pp.type_pp("Node") as tpp:
 479             tpp.field("graph", ...)
 480             tpp.field("id", _NodeSimpleReprWrapper(self))
 481             tpp.try_field("root",
 482                           lambda: _NodeSimpleReprWrapper(self.root),
 483                           BudgetGraphLoopError)
 484             parent = f"#{self.parent_id}" if self.parent_id is not None else None
 485             tpp.field("parent", parent)
 486             tpp.field("budget_excluding_subtasks",
 487                       self.budget_excluding_subtasks)
 488             tpp.field("budget_including_subtasks",
 489                       self.budget_including_subtasks)
 490             tpp.field("fixed_budget_excluding_subtasks",
 491                       self.fixed_budget_excluding_subtasks)
 492             tpp.field("fixed_budget_including_subtasks",
 493                       self.fixed_budget_including_subtasks)
 494             tpp.field("milestone_str", self.milestone_str)
 495             tpp.try_field("milestone", lambda: self.milestone,
 496                           BudgetGraphBaseError)
 497             immediate_children = [_NodeSimpleReprWrapper(i)
 498                                   for i in self.immediate_children]
 499             tpp.field("immediate_children", immediate_children)
 500             tpp.try_field("payments",
 501                           lambda: list(self.payments.values()),
 502                           BudgetGraphBaseError)
 503             try:
 504                 status = repr(self.status)
 505             except BudgetGraphBaseError:
 506                 status = f"<unknown status: {self.bug.status!r}>"
 507             tpp.field("status", status)
 508             try:
 509                 assignee = f"Person<{self.assignee.identifier!r}>"
 510             except BudgetGraphBaseError:
 511                 assignee = f"<unknown assignee: {self.bug.assigned_to!r}>"
 512             tpp.field("assignee", assignee)
 513             tpp.try_field("resolved_payments",
 514                           lambda: self.resolved_payments,
 515                           BudgetGraphBaseError)
 516             tpp.try_field("payment_summaries",
 517                           lambda: self.payment_summaries,
 518                           BudgetGraphBaseError)
 519
 520     def __repr__(self):
 521         try:
 522             root = _NodeSimpleReprWrapper(self.root)
 523         except BudgetGraphLoopError:
 524             root = "<loop error>"
 525         try:
 526             milestone = repr(self.milestone)
 527         except BudgetGraphBaseError:
 528             milestone = "<unknown milestone>"
 529         try:
 530             status = repr(self.status)
 531         except BudgetGraphBaseError:
 532             status = f"<unknown status: {self.bug.status!r}>"
 533         try:
 534             assignee = f"Person<{self.assignee.identifier!r}>"
 535         except BudgetGraphBaseError:
 536             assignee = f"<unknown assignee: {self.bug.assigned_to!r}>"
 537         immediate_children = []
 538         for i in self.immediate_children:
 539             immediate_children.append(_NodeSimpleReprWrapper(i))
 540         immediate_children.sort()
 541         parent = f"#{self.parent_id}" if self.parent_id is not None else None
 542         payments = list(self.payments.values())
 543         resolved_payments = self.resolved_payments
 544         payment_summaries = self.payment_summaries
 545         return (f"Node(graph=..., "
 546                 f"id={_NodeSimpleReprWrapper(self)}, "
 547                 f"root={root}, "
 548                 f"parent={parent}, "
 549                 f"budget_excluding_subtasks={self.budget_excluding_subtasks}, "
 550                 f"budget_including_subtasks={self.budget_including_subtasks}, "
 551                 f"fixed_budget_excluding_subtasks={self.fixed_budget_excluding_subtasks}, "
 552                 f"fixed_budget_including_subtasks={self.fixed_budget_including_subtasks}, "
 553                 f"milestone_str={self.milestone_str!r}, "
 554                 f"milestone={milestone}, "
 555                 f"immediate_children={immediate_children!r}, "
 556                 f"payments={payments!r}, "
 557                 f"status={status}, "
 558                 f"assignee={assignee}, "
 559                 f"resolved_payments={resolved_payments!r}, "
 560                 f"payment_summaries={payment_summaries!r})")
 561
 562
 563 class BudgetGraphError(BudgetGraphBaseError):
 564     def __init__(self, bug_id: int, root_bug_id: int):
 565         self.bug_id = bug_id
 566         self.root_bug_id = root_bug_id
 567
 568
 569 class BudgetGraphMoneyWithNoMilestone(BudgetGraphError):
 570     def __str__(self):
 571         return (f"Bug assigned money but without"
 572                 f" any assigned milestone: #{self.bug_id}")
 573
 574
 575 class BudgetGraphMilestoneMismatch(BudgetGraphError):
 576     def __str__(self):
 577         return (f"Bug's assigned milestone doesn't match the milestone "
 578                 f"assigned to the root bug: descendant bug"
 579                 f" #{self.bug_id}, root bug"
 580                 f" #{self.root_bug_id}")
 581
 582
 583 class BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(BudgetGraphError):
 584     def __init__(self, bug_id: int, root_bug_id: int,
 585                  expected_budget_excluding_subtasks: Money):
 586         super().__init__(bug_id, root_bug_id)
 587         self.expected_budget_excluding_subtasks = \
 588             expected_budget_excluding_subtasks
 589
 590     def __str__(self):
 591         return (f"Budget assigned to task excluding subtasks "
 592                 f"(cf_budget field) doesn't match calculated value: "
 593                 f"bug #{self.bug_id}, calculated value"
 594                 f" {self.expected_budget_excluding_subtasks}")
 595
 596
 597 class BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(BudgetGraphError):
 598     def __init__(self, bug_id: int, root_bug_id: int,
 599                  expected_budget_including_subtasks: Money):
 600         super().__init__(bug_id, root_bug_id)
 601         self.expected_budget_including_subtasks = \
 602             expected_budget_including_subtasks
 603
 604     def __str__(self):
 605         return (f"Budget assigned to task including subtasks "
 606                 f"(cf_total_budget field) doesn't match calculated value: "
 607                 f"bug #{self.bug_id}, calculated value"
 608                 f" {self.expected_budget_including_subtasks}")
 609
 610
 611 class BudgetGraphNegativeMoney(BudgetGraphError):
 612     def __str__(self):
 613         return (f"Budget assigned to task is less than zero: "
 614                 f"bug #{self.bug_id}")
 615
 616
 617 class BudgetGraphPayeesMoneyMismatch(BudgetGraphError):
 618     def __init__(self, bug_id: int, root_bug_id: int, payees_total: Money,
 619                  expected_payees_total: Money):
 620         super().__init__(bug_id, root_bug_id)
 621         self.payees_total = payees_total
 622         self.expected_payees_total = expected_payees_total
 623
 624     def __str__(self):
 625         return (f"Total budget assigned to payees (cf_payees_list) doesn't "
 626                 f"match expected value: bug #{self.bug_id}, calculated total "
 627                 f"{self.payees_total}, expected value "
 628                 f"{self.expected_payees_total}")
 629
 630
 631 class BudgetGraphNegativePayeeMoney(BudgetGraphError):
 632     def __init__(self, bug_id: int, root_bug_id: int, payee_key: str):
 633         super().__init__(bug_id, root_bug_id)
 634         self.payee_key = payee_key
 635
 636     def __str__(self):
 637         return (f"Budget assigned to payee for task is less than zero: "
 638                 f"bug #{self.bug_id}, payee {self.payee_key!r}")
 639
 640
 641 class BudgetGraphIncorrectRootForMilestone(BudgetGraphError):
 642     def __init__(self, bug_id: int, milestone: str, milestone_canonical_bug_id: int):
 643         super().__init__(bug_id, bug_id)
 644         self.milestone = milestone
 645         self.milestone_canonical_bug_id = milestone_canonical_bug_id
 646
 647     def __str__(self):
 648         return (f"Bug #{self.bug_id} is not the canonical root bug for "
 649                 f"assigned milestone {self.milestone!r} but has no parent "
 650                 f"bug set: the milestone's canonical root bug is "
 651                 f"#{self.milestone_canonical_bug_id}")
 652
 653
 654 class BudgetGraph:
 655     nodes: Dict[int, Node]
 656
 657     def __init__(self, bugs: Iterable[Bug], config: Config):
 658         self.nodes = {}
 659         self.config = config
 660         for bug in bugs:
 661             self.nodes[bug.id] = Node(self, bug)
 662         for node in self.nodes.values():
 663             if node.parent is None:
 664                 continue
 665             node.parent.immediate_children.add(node)
 666         # useful debug prints
 667         # for bug in bugs:
 668         #    node = self.nodes[bug.id]
 669         #    print ("bug added", bug.id, node, node.parent.immediate_children)
 670
 671     @cached_property
 672     def roots(self) -> OrderedSet[Node]:
 673         roots = OrderedSet()
 674         for node in self.nodes.values():
 675             # calling .root also checks for loop errors
 676             root = node.root
 677             roots.add(root)
 678         return roots
 679
 680     def _get_node_errors(self, root: Node, node: Node,
 681                          errors: List[BudgetGraphBaseError]):
 682         if node.milestone_str is None:
 683             if node.budget_including_subtasks != 0 \
 684                     or node.budget_excluding_subtasks != 0:
 685                 errors.append(BudgetGraphMoneyWithNoMilestone(
 686                     node.bug.id, root.bug.id))
 687
 688         try:
 689             # check for milestone errors
 690             node.milestone
 691             if root == node and node.milestone is not None \
 692                     and node.milestone.canonical_bug_id != node.bug.id:
 693                 if node.budget_including_subtasks != 0 \
 694                         or node.budget_excluding_subtasks != 0:
 695                     errors.append(BudgetGraphIncorrectRootForMilestone(
 696                         node.bug.id, node.milestone.identifier,
 697                         node.milestone.canonical_bug_id
 698                     ))
 699         except BudgetGraphBaseError as e:
 700             errors.append(e)
 701
 702         try:
 703             # check for status errors
 704             node.status
 705         except BudgetGraphBaseError as e:
 706             errors.append(e)
 707
 708         try:
 709             # check for assignee errors
 710             node.assignee
 711         except BudgetGraphBaseError as e:
 712             errors.append(e)
 713
 714         if node.milestone_str != root.milestone_str:
 715             errors.append(BudgetGraphMilestoneMismatch(
 716                 node.bug.id, root.bug.id))
 717
 718         if node.budget_excluding_subtasks < 0 \
 719                 or node.budget_including_subtasks < 0:
 720             errors.append(BudgetGraphNegativeMoney(
 721                 node.bug.id, root.bug.id))
 722
 723         childlist = []
 724         subtasks_total = Money(0)
 725         for child in node.immediate_children:
 726             subtasks_total += child.fixed_budget_including_subtasks
 727             childlist.append(child.bug.id)
 728         # useful debug prints
 729         # print ("subtask total", node.bug.id, root.bug.id, subtasks_total,
 730         #                        childlist)
 731
 732         payees_total = Money(0)
 733         payee_payments: Dict[Person, List[Payment]] = {}
 734         for payment in node.payments.values():
 735             if payment.amount < 0:
 736                 errors.append(BudgetGraphNegativePayeeMoney(
 737                     node.bug.id, root.bug.id, payment.payee_key))
 738             payees_total += payment.amount
 739             try:
 740                 # check for payee errors
 741                 payment.payee
 742                 previous_payment = payee_payments.get(payment.payee)
 743                 if previous_payment is not None:
 744                     payee_payments[payment.payee].append(payment)
 745                 else:
 746                     payee_payments[payment.payee] = [payment]
 747             except BudgetGraphBaseError as e:
 748                 errors.append(e)
 749
 750         def set_including_from_excluding_and_error():
 751             node.fixed_budget_including_subtasks = \
 752                 node.budget_excluding_subtasks + subtasks_total
 753             errors.append(
 754                 BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(
 755                     node.bug.id, root.bug.id,
 756                     node.fixed_budget_including_subtasks))
 757
 758         def set_including_from_payees_and_error():
 759             node.fixed_budget_including_subtasks = \
 760                 payees_total + subtasks_total
 761             errors.append(
 762                 BudgetGraphMoneyMismatchForBudgetIncludingSubtasks(
 763                     node.bug.id, root.bug.id,
 764                     node.fixed_budget_including_subtasks))
 765
 766         def set_excluding_from_including_and_error():
 767             v = node.budget_including_subtasks - subtasks_total
 768             if v < 0:
 769                 return set_including_from_excluding_and_error()
 770             node.fixed_budget_excluding_subtasks = v
 771             errors.append(
 772                 BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(
 773                     node.bug.id, root.bug.id,
 774                     node.fixed_budget_excluding_subtasks))
 775
 776         def set_excluding_from_payees_and_error():
 777             node.fixed_budget_excluding_subtasks = \
 778                 payees_total
 779             errors.append(
 780                 BudgetGraphMoneyMismatchForBudgetExcludingSubtasks(
 781                     node.bug.id, root.bug.id,
 782                     node.fixed_budget_excluding_subtasks))
 783
 784         def set_payees_from_including_and_error():
 785             fixed_payees_total = \
 786                 node.budget_including_subtasks - subtasks_total
 787             errors.append(BudgetGraphPayeesMoneyMismatch(
 788                 node.bug.id, root.bug.id, payees_total, fixed_payees_total))
 789
 790         def set_payees_from_excluding_and_error():
 791             fixed_payees_total = \
 792                 node.budget_excluding_subtasks
 793             errors.append(BudgetGraphPayeesMoneyMismatch(
 794                 node.bug.id, root.bug.id, payees_total, fixed_payees_total))
 795
 796         payees_matches_including = \
 797             node.budget_including_subtasks - subtasks_total == payees_total
 798         payees_matches_excluding = \
 799             node.budget_excluding_subtasks == payees_total
 800         including_matches_excluding = \
 801             node.budget_including_subtasks - subtasks_total \
 802             == node.budget_excluding_subtasks
 803
 804         if payees_matches_including \
 805                 and payees_matches_excluding \
 806                 and including_matches_excluding:
 807             pass  # no error
 808         elif payees_matches_including:
 809             # can't have 2 match without all 3 matching
 810             assert not payees_matches_excluding
 811             assert not including_matches_excluding
 812             if node.budget_including_subtasks == 0 and len(node.payments) == 0:
 813                 set_including_from_excluding_and_error()
 814             else:
 815                 set_excluding_from_including_and_error()
 816         elif payees_matches_excluding:
 817             # can't have 2 match without all 3 matching
 818             assert not payees_matches_including
 819             assert not including_matches_excluding
 820             if node.budget_excluding_subtasks == 0 and len(node.payments) == 0:
 821                 if node.budget_including_subtasks == 0:
 822                     set_including_from_excluding_and_error()
 823                 else:
 824                     set_excluding_from_including_and_error()
 825             else:
 826                 set_including_from_excluding_and_error()
 827         elif including_matches_excluding:
 828             # can't have 2 match without all 3 matching
 829             assert not payees_matches_including
 830             assert not payees_matches_excluding
 831             if len(node.payments) == 0:
 832                 pass  # no error -- payees is just not set
 833             elif node.budget_excluding_subtasks == 0 \
 834                     and node.budget_including_subtasks == 0:
 835                 set_excluding_from_payees_and_error()
 836                 set_including_from_payees_and_error()
 837             else:
 838                 set_payees_from_excluding_and_error()
 839         else:
 840             # nothing matches
 841             if len(node.payments) == 0:
 842                 # payees unset -- don't need to set payees
 843                 if node.budget_including_subtasks == 0:
 844                     set_including_from_excluding_and_error()
 845                 else:
 846                     set_excluding_from_including_and_error()
 847             elif node.budget_excluding_subtasks == 0 \
 848                     and node.budget_including_subtasks == 0:
 849                 set_excluding_from_payees_and_error()
 850                 set_including_from_payees_and_error()
 851             elif node.budget_excluding_subtasks == 0:
 852                 set_excluding_from_including_and_error()
 853                 set_payees_from_including_and_error()
 854             elif node.budget_including_subtasks == 0:
 855                 set_including_from_excluding_and_error()
 856                 set_payees_from_excluding_and_error()
 857             else:
 858                 set_including_from_excluding_and_error()
 859                 set_payees_from_excluding_and_error()
 860
 861     def get_errors(self) -> List[BudgetGraphBaseError]:
 862         errors = []
 863         try:
 864             roots = self.roots
 865         except BudgetGraphBaseError as e:
 866             errors.append(e)
 867             return errors
 868
 869         for root in roots:
 870             try:
 871                 for child in reversed(list(root.children_breadth_first())):
 872                     try:
 873                         self._get_node_errors(root, child, errors)
 874                     except BudgetGraphBaseError as e:
 875                         errors.append(e)
 876                 self._get_node_errors(root, root, errors)
 877             except BudgetGraphBaseError as e:
 878                 errors.append(e)
 879         return errors
 880
 881     @cached_property
 882     def assigned_nodes(self) -> Dict[Person, List[Node]]:
 883         retval: Dict[Person, List[Node]]
 884         retval = {person: [] for person in self.config.people.values()}
 885         for node in self.nodes.values():
 886             retval[node.assignee].append(node)
 887         return retval
 888
 889     @cached_property
 890     def assigned_nodes_for_milestones(self) -> Dict[Milestone, List[Node]]:
 891         retval: Dict[Milestone, List[Node]]
 892         retval = {milestone: []
 893                   for milestone in self.config.milestones.values()}
 894         for node in self.nodes.values():
 895             if node.milestone is not None:
 896                 retval[node.milestone].append(node)
 897         return retval
 898
 899     @cached_property
 900     def milestone_payments(self) -> Dict[Milestone, List[Payment]]:
 901         retval: Dict[Milestone, List[Payment]] = {
 902             milestone: [] for milestone in self.config.milestones.values()
 903         }
 904         for node in self.nodes.values():
 905             if node.milestone is not None:
 906                 retval[node.milestone].extend(node.payments.values())
 907         return retval
 908
 909     @cached_property
 910     def payments(self) -> Dict[Person, Dict[Milestone, List[Payment]]]:
 911         retval: Dict[Person, Dict[Milestone, List[Payment]]] = {
 912             person: {
 913                 milestone: []
 914                 for milestone in self.config.milestones.values()
 915             }
 916             for person in self.config.people.values()
 917         }
 918         for node in self.nodes.values():
 919             if node.milestone is not None:
 920                 for payment in node.payments.values():
 921                     retval[payment.payee][node.milestone].append(payment)
 922         return retval
 923
 924     @cached_property
 925     def milestone_people(self) -> Dict[Milestone, OrderedSet[Person]]:
 926         """get a list of people associated with each milestone
 927         """
 928         payments = list(self.payments)  # just activate the payments
 929         retval = {}
 930         for milestone in self.milestone_payments.keys():
 931             retval[milestone] = OrderedSet()
 932         for milestone, payments in self.milestone_payments.items():
 933             for payment in payments:
 934                 retval[milestone].add(payment.payee)
 935         return retval
 936
 937     def __pretty_print__(self, pp: PrettyPrinter):
 938         with pp.type_pp("BudgetGraph") as tpp:
 939             tpp.field("nodes", self.nodes)
 940             tpp.try_field("roots",
 941                           lambda: [_NodeSimpleReprWrapper(i)
 942                                    for i in self.roots],
 943                           BudgetGraphBaseError)
 944             tpp.try_field("assigned_nodes",
 945                           lambda: {
 946                               person: [
 947                                   _NodeSimpleReprWrapper(node)
 948                                   for node in nodes
 949                               ]
 950                               for person, nodes in self.assigned_nodes.items()
 951                           },
 952                           BudgetGraphBaseError)
 953             tpp.try_field("assigned_nodes_for_milestones",
 954                           lambda: {
 955                               milestone: [
 956                                   _NodeSimpleReprWrapper(node)
 957                                   for node in nodes
 958                               ]
 959                               for milestone, nodes in self.assigned_nodes_for_milestones.items()
 960                           },
 961                           BudgetGraphBaseError)
 962             tpp.try_field("payments",
 963                           lambda: self.payments, BudgetGraphBaseError)
 964             tpp.try_field("milestone_people",
 965                           lambda: self.milestone_people,
 966                           BudgetGraphBaseError)
 967
 968     def __repr__(self):
 969         nodes = [*self.nodes.values()]
 970
 971         def repr_or_failed(f: Callable[[], Any]) -> str:
 972             try:
 973                 return repr(f())
 974             except BudgetGraphBaseError:
 975                 return "<failed>"
 976
 977         try:
 978             roots = [_NodeSimpleReprWrapper(i) for i in self.roots]
 979             roots.sort()
 980             roots_str = repr(roots)
 981         except BudgetGraphBaseError:
 982             roots_str = "<failed>"
 983         assigned_nodes = repr_or_failed(lambda: {
 984             person: [
 985                 _NodeSimpleReprWrapper(node)
 986                 for node in nodes
 987             ]
 988             for person, nodes in self.assigned_nodes.items()
 989         })
 990         assigned_nodes_for_milestones = repr_or_failed(lambda: {
 991             milestone: [
 992                 _NodeSimpleReprWrapper(node)
 993                 for node in nodes
 994             ]
 995             for milestone, nodes in self.assigned_nodes_for_milestones.items()
 996         })
 997         milestone_payments = repr_or_failed(lambda: self.milestone_payments)
 998         payments = repr_or_failed(lambda: self.payments)
 999         milestone_people = repr_or_failed(lambda: self.milestone_people)
1000         return (f"BudgetGraph{{nodes={nodes!r}, "
1001                 f"roots={roots}, "
1002                 f"assigned_nodes={assigned_nodes}, "
1003                 f"assigned_nodes_for_milestones={assigned_nodes_for_milestones}, "
1004                 f"milestone_payments={milestone_payments}, "
1005                 f"payments={payments}, "
1006                 f"milestone_people={milestone_people}}}")