Allow decimals in F1 files

gisce · Oct 14, 2024 · 3dfe54c · 3dfe54c
1 parent 5323842
commit 3dfe54c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/mesures/f1.py b/mesures/f1.py
@@ -9,7 +9,7 @@
 
 
 class F1(object):
-    def __init__(self, data, distributor=None, compression='bz2', version=0):
+    def __init__(self, data, distributor=None, compression='bz2', allow_decimals=False, version=0):
         """
         :param data: list of dicts or absolute file_path
         :param distributor: str distributor REE code
@@ -24,6 +24,7 @@ def __init__(self, data, distributor=None, compression='bz2', version=0):
         self.version = version
         self.distributor = distributor
         self.default_compression = compression
+        self.allow_decimals = allow_decimals
 
     def __repr__(self):
         return "{}: {} kWh".format(self.filename, self.total)
@@ -129,6 +130,11 @@ def reader(self, filepath):
 
         df['res'] = 0
         df['res2'] = 0
+
+        if not self.allow_decimals:
+            for key in ['ai', 'ae', 'r1', 'r2', 'r3', 'r4']:
+                df[key] = df[key].astype('int')
+
         df = df[self.columns]
         return df
 

diff --git a/mesures/f1qh.py b/mesures/f1qh.py
@@ -8,13 +8,14 @@
 
 
 class F1QH(F1):
-    def __init__(self, data, distributor=None, compression='bz2', version=0):
+    def __init__(self, data, distributor=None, compression='bz2', allow_decimals=False, version=0):
         """
         :param data: list of dicts or absolute file_path
         :param distributor: str distributor REE code
         :param compression: 'bz2', 'gz'... OR False otherwise
         """
-        super(F1QH, self).__init__(data=data, distributor=distributor, compression=compression, version=version)
+        super(F1QH, self).__init__(data=data, distributor=distributor, compression=compression,
+                                   allow_decimals=allow_decimals, version=version)
         self.prefix = 'F1QH'
 
     def reader(self, filepath):
@@ -49,6 +50,11 @@ def reader(self, filepath):
 
         df['res'] = 0
         df['res2'] = 0
+
+        if not self.allow_decimals:
+            for key in ['ai', 'ae', 'r1', 'r2', 'r3', 'r4']:
+                df[key] = df[key].astype('int')
+
         df = df[self.columns]
         return df