From 5df10d16b371ec37a144bdcb847290acc2f9b2a3 Mon Sep 17 00:00:00 2001 From: Josh Renaud Date: Wed, 6 Oct 2021 14:31:50 -0500 Subject: [PATCH] Potential revision for mean.py to add support for TimeDelta. https://github.com/wireservice/agate/issues/761 --- agate/aggregations/mean.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/agate/aggregations/mean.py b/agate/aggregations/mean.py index 689b674f..ac0de983 100644 --- a/agate/aggregations/mean.py +++ b/agate/aggregations/mean.py @@ -3,7 +3,7 @@ from agate.aggregations.base import Aggregation from agate.aggregations.has_nulls import HasNulls from agate.aggregations.sum import Sum -from agate.data_types import Number +from agate.data_types import Number, TimeDelta from agate.exceptions import DataTypeError from agate.warns import warn_null_calculation @@ -20,13 +20,16 @@ def __init__(self, column_name): self._sum = Sum(column_name) def get_aggregate_data_type(self, table): - return Number() + column = table.columns[self._column_name] + + if isinstance(column.data_type, (Number, TimeDelta)): + return column.data_type def validate(self, table): column = table.columns[self._column_name] - if not isinstance(column.data_type, Number): - raise DataTypeError('Mean can only be applied to columns containing Number data.') + if not isinstance(column.data_type, (Number, TimeDelta)): + raise DataTypeError('Sum can only be applied to columns containing Number or TimeDelta data.') has_nulls = HasNulls(self._column_name).run(table)