diff --git a/conservancy_beancount/plugin/core.py b/conservancy_beancount/plugin/core.py index befa40c..d3cb6a8 100644 --- a/conservancy_beancount/plugin/core.py +++ b/conservancy_beancount/plugin/core.py @@ -36,6 +36,36 @@ class _GenericRange: return self.start <= item < self.stop +class MetadataEnum: + def __init__(self, key, standard_values, aliases_map): + self.key = key + self._stdvalues = frozenset(standard_values) + self._aliases = dict(aliases_map) + self._aliases.update((v, v) for v in standard_values) + assert self._stdvalues == set(self._aliases.values()) + + def __repr__(self): + return "{}<{}>".format(type(self).__name__, self.key) + + def __contains__(self, key): + return key in self._aliases + + def __getitem__(self, key): + return self._aliases[key] + + def __iter__(self): + return iter(self._stdvalues) + + def get(self, key, default_key=None): + try: + return self[key] + except KeyError: + if default_key is None: + return None + else: + return self[default_key] + + class PostingChecker: ACCOUNTS = ('',) TXN_DATE_RANGE = _GenericRange(DEFAULT_START_DATE, DEFAULT_STOP_DATE) @@ -76,7 +106,7 @@ class PostingChecker: errors.append(error) else: try: - set_value = self.VALUES_ENUM[source_value].value + set_value = self.VALUES_ENUM[source_value] except KeyError: errors.append(errormod.InvalidMetadataError( txn, post, self.METADATA_KEY, source_value, diff --git a/conservancy_beancount/plugin/meta_expense_allocation.py b/conservancy_beancount/plugin/meta_expense_allocation.py index e1e732e..603b2d7 100644 --- a/conservancy_beancount/plugin/meta_expense_allocation.py +++ b/conservancy_beancount/plugin/meta_expense_allocation.py @@ -14,27 +14,23 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -import enum - from . import core -class ExpenseAllocations(enum.Enum): - administration = 'administration' - fundraising = 'fundraising' - program = 'program' - - admin = administration - - class MetaExpenseAllocation(core.PostingChecker): ACCOUNTS = ('Expenses:',) METADATA_KEY = 'expenseAllocation' - VALUES_ENUM = ExpenseAllocations + VALUES_ENUM = core.MetadataEnum(METADATA_KEY, { + 'administration', + 'fundraising', + 'program', + }, { + 'admin': 'administration', + }) DEFAULT_VALUES = { - 'Expenses:Services:Accounting': VALUES_ENUM.administration, - 'Expenses:Services:Administration': VALUES_ENUM.administration, - 'Expenses:Services:Fundraising': VALUES_ENUM.fundraising, + 'Expenses:Services:Accounting': VALUES_ENUM['administration'], + 'Expenses:Services:Administration': VALUES_ENUM['administration'], + 'Expenses:Services:Fundraising': VALUES_ENUM['fundraising'], } def _default_value(self, txn, post): - return self.DEFAULT_VALUES.get(post.account, self.VALUES_ENUM.program).value + return self.DEFAULT_VALUES.get(post.account, 'program')