diff --git a/mesures/f5d.py b/mesures/f5d.py index 4c8f56e..2380076 100644 --- a/mesures/f5d.py +++ b/mesures/f5d.py @@ -9,20 +9,29 @@ TYPES = DTYPES.copy() TYPES.update({'factura': 'category'}) - +ENERGY_MAGNS = ['ai', 'ae', 'r1', 'r2', 'r3', 'r4'] +CNMC_ENERGY_MAGNS = ENERGY_MAGNS + ['ai_fix', 'ae_fix'] class F5D(F5): - def __init__(self, data, distributor=None, comer=None, compression='bz2', columns=COLUMNS, dtypes=TYPES, version=0): + def __init__(self, data, file_format='REE', distributor=None, comer=None, compression='bz2', columns=COLUMNS, dtypes=TYPES, version=0): """ :param data: list of dicts or absolute file_path + :param file_format: str format to generate :param distributor: str distributor REE code :param comer: str comer REE code :param compression: 'bz2', 'gz'... OR False otherwise """ + + self.file_format = file_format + if self.file_format == 'CNMC': + columns = columns + CNMC_ENERGY_MAGNS + self.columns = columns + super(F5D, self).__init__(data, distributor=distributor, comer=comer, compression=compression, columns=columns, dtypes=dtypes, version=version) self.prefix = 'F5D' + @property def filename(self): filename = "{prefix}_{distributor}_{comer}_{timestamp}.{version}".format( @@ -62,14 +71,21 @@ def reader(self, filepath): if 'factura' not in df: df['factura'] = 'F0000000000' + agregates = {'ai': 'sum', 'ae': 'sum', 'r1': 'sum', 'r2': 'sum', 'r3': 'sum', 'r4': 'sum'} + if self.file_format == 'CNMC': + agregates.update({'ai_fix': 'sum', 'ae_fix': 'sum'}) + df = df.groupby(['cups', 'timestamp', 'season', 'firmeza', 'method', 'factura']).aggregate( - {'ai': 'sum', 'ae': 'sum', 'r1': 'sum', 'r2': 'sum', 'r3': 'sum', 'r4': 'sum'} - ).reset_index() + agregates).reset_index() if isinstance(filepath, list): df['timestamp'] = df['timestamp'].apply(lambda x: x.strftime('%Y/%m/%d %H:%M')) - for key in ['ai', 'ae', 'r1', 'r2', 'r3', 'r4']: + magnituds = ENERGY_MAGNS + if self.file_format == 'CNMC': + magnituds = CNMC_ENERGY_MAGNS + + for key in magnituds: if key not in df: df[key] = 0 df[key] = df[key].astype('int32')