first approach: Adapt to new CNMC F5D format

gisce · Dec 11, 2024 · 60e57c8 · 60e57c8
1 parent 6b974e9
commit 60e57c8
Showing 1 changed file with 21 additions and 5 deletions.
diff --git 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')