diff --git a/mesures/f1.py b/mesures/f1.py index 08f65bf..f062bbc 100644 --- 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 index 3c25f00..02391c6 100644 --- 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