src/budget_sync/write_budget_csv.py

   1 import csv
   2 from enum import Enum, auto
   3 from io import StringIO
   4 from typing import Any, Callable, Dict, List, Optional
   5 from budget_sync.budget_graph import BudgetGraph, Node, PayeeState, PaymentSummary
   6 from pathlib import Path
   7 from budget_sync.config import Milestone
   8 from budget_sync.money import Money
   9 from budget_sync.write_budget_markdown import markdown_escape
  10
  11
  12 def _budget_csv_row(budget_graph: BudgetGraph, milestone: Milestone, node: Optional[Node]) -> Dict[str, str]:
  13     row_fns: Dict[str, Callable[[Node], Any]] = {
  14         'bug_id': lambda node: node.bug.id,
  15         'excl_subtasks': lambda node: node.budget_excluding_subtasks,
  16         'inc_subtasks': lambda node: node.budget_including_subtasks,
  17         'fixed_excl_subtasks': lambda node: node.fixed_budget_excluding_subtasks,
  18         'fixed_inc_subtasks': lambda node: node.fixed_budget_including_subtasks,
  19         'req_excl_subtasks': lambda node: node.submitted_excluding_subtasks,
  20         'paid_excl_subtasks': lambda node: node.paid_excluding_subtasks,
  21     }
  22     milestone_people = budget_graph.milestone_people[milestone]
  23
  24     def handle_person(person):
  25         # need a nested function in order to create a new person variable
  26         # for this iteration that can be bound to the lambdas
  27         id = person.identifier
  28         row_fns.update({
  29             id + " (planned amt)": lambda node: node.payment_summaries[person].total,
  30             id + " (req amt)": lambda node: node.payment_summaries[person].total_submitted,
  31             id + " (req date)": lambda node: node.payment_summaries[person].submitted_date,
  32             id + " (paid amt)": lambda node: node.payment_summaries[person].total_paid,
  33             id + " (paid date)": lambda node: node.payment_summaries[person].paid_date,
  34         })
  35     for person in milestone_people:
  36         handle_person(person)
  37     row = {k: "" for k in row_fns.keys()}
  38     if node is None:
  39         return row
  40     for k, fn in row_fns.items():
  41         try:
  42             v = fn(node)
  43         except KeyError:
  44             continue
  45         if v is not None:
  46             row[k] = str(v)
  47     return row
  48
  49
  50 def _budget_csv_for_milestone(budget_graph: BudgetGraph, milestone: Milestone) -> str:
  51     with StringIO() as string_io:
  52         writer = csv.DictWriter(
  53             string_io,
  54             _budget_csv_row(budget_graph, milestone, None).keys(),
  55             lineterminator="\n")
  56         writer.writeheader()
  57         for node in budget_graph.assigned_nodes_for_milestones[milestone]:
  58             # skip uninteresting nodes
  59             if len(node.payments) == 0 \
  60                     and node.budget_excluding_subtasks == 0 \
  61                     and node.budget_including_subtasks == 0:
  62                 continue
  63             row = _budget_csv_row(budget_graph, milestone, node)
  64             writer.writerow(row)
  65         return string_io.getvalue()
  66
  67
  68 def write_budget_csv(budget_graph: BudgetGraph,
  69                      output_dir: Path):
  70     output_dir.mkdir(parents=True, exist_ok=True)
  71     milestones = budget_graph.config.milestones
  72     csv_paths: Dict[Milestone, Path] = {}
  73     for milestone in milestones.values():
  74         csv_text = _budget_csv_for_milestone(budget_graph, milestone)
  75         csv_paths[milestone] = output_dir.joinpath(
  76             f"{milestone.identifier}.csv")
  77         csv_paths[milestone].write_text(csv_text, encoding="utf-8")
  78
  79     markdown_text = "\n".join(f"# {markdown_escape(milestone.identifier)}\n"
  80                               "\n"
  81                               f"[[!table format=csv file=\"{path!s}\"]]"
  82                               for milestone, path in csv_paths.items())
  83     output_dir.joinpath("csvs.mdwn").write_text(
  84         markdown_text, encoding="utf-8")