"""Common attributes for totals."""

from pynamodb.attributes import NumberAttribute

from sosmodels import consts


# All PynamoDB attributes must belong to particular Attribute container (
# present in __dict__ attribute) so we have to inject them to each class
def with_overall_totals(cls):
    """Append AttributeContainer with stream totals fields."""
    for total_metric in consts.TOTAL_METRICS:
        setattr(cls, total_metric, NumberAttribute(default=0))
    return cls