108 lines
3.4 KiB
Python
108 lines
3.4 KiB
Python
"""Enhanced Beancount data structures for Conservancy"""
|
|
# Copyright © 2020 Brett Smith
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
import collections.abc
|
|
|
|
from beancount.core import account as bc_account
|
|
|
|
from typing import (
|
|
Iterable,
|
|
Iterator,
|
|
MutableMapping,
|
|
Optional,
|
|
)
|
|
|
|
from .beancount_types import (
|
|
MetaKey,
|
|
MetaValue,
|
|
Posting as BasePosting,
|
|
Transaction,
|
|
)
|
|
|
|
class Account(str):
|
|
SEP = bc_account.sep
|
|
|
|
def is_under(self, acct_s: str) -> bool:
|
|
return self.startswith(acct_s) and (
|
|
acct_s.endswith(self.SEP)
|
|
or self == acct_s
|
|
or self[len(acct_s)] == self.SEP
|
|
)
|
|
|
|
|
|
class PostingMeta(collections.abc.MutableMapping):
|
|
def __init__(self,
|
|
txn: Transaction,
|
|
index: int,
|
|
post: Optional[BasePosting]=None,
|
|
) -> None:
|
|
if post is None:
|
|
post = txn.postings[index]
|
|
self.txn = txn
|
|
self.index = index
|
|
self.post = post
|
|
|
|
def __iter__(self) -> Iterator[MetaKey]:
|
|
keys: Iterable[MetaKey]
|
|
if self.post.meta is None:
|
|
keys = self.txn.meta.keys()
|
|
else:
|
|
keys = frozenset(self.post.meta.keys()).union(self.txn.meta.keys())
|
|
return iter(keys)
|
|
|
|
def __len__(self) -> int:
|
|
return sum(1 for _ in self)
|
|
|
|
def __getitem__(self, key: MetaKey) -> MetaValue:
|
|
if self.post.meta:
|
|
try:
|
|
return self.post.meta[key]
|
|
except KeyError:
|
|
pass
|
|
return self.txn.meta[key]
|
|
|
|
def __setitem__(self, key: MetaKey, value: MetaValue) -> None:
|
|
if self.post.meta is None:
|
|
self.post = self.post._replace(meta={key: value})
|
|
self.txn.postings[self.index] = self.post
|
|
else:
|
|
self.post.meta[key] = value
|
|
|
|
def __delitem__(self, key: MetaKey) -> None:
|
|
if self.post.meta is None:
|
|
raise KeyError(key)
|
|
else:
|
|
del self.post.meta[key]
|
|
|
|
|
|
class Posting(BasePosting):
|
|
account: Account
|
|
# mypy correctly complains that our MutableMapping is not compatible with
|
|
# Beancount's meta type declaration of Optional[Dict]. IMO this is a case
|
|
# of Beancount's type declaration being a smidge too specific: I think it
|
|
# would be very unusual for code to actually require a dict over a more
|
|
# generic mapping. If it did, this would work fine.
|
|
meta: MutableMapping[MetaKey, MetaValue] # type:ignore[assignment]
|
|
|
|
|
|
def iter_postings(txn: Transaction) -> Iterator[Posting]:
|
|
for index, source in enumerate(txn.postings):
|
|
yield Posting(
|
|
Account(source.account),
|
|
*source[1:5],
|
|
# see rationale above about Posting.meta
|
|
PostingMeta(txn, index, source), # type:ignore[arg-type]
|
|
)
|