Merge pull request #3 from Aymane11/fix-double-log

Fix double logging issue

anonymize/__main__.py

@@ -1,4 +1,3 @@
-import sys
 from anonymize import load_config, Config
 
 from loguru import logger
@@ -9,21 +8,13 @@
 def main(config: Config):
     dfs: list[pl.LazyFrame] = []
     for data in config.source:
-        if len(dfs) > 0:
-            # In case of database, no need to log every time
-            logger.remove()
         for rule in config.rules:
             if rule.column not in data.columns:
                 logger.warning(f"Column {rule.column} not found in the dataset. Skipping.")
                 continue
             data = rule.apply(data)
         dfs.append(data)
 
-    logger.add(
-        sys.stdout,
-        colorize=True,
-        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <level>{message}</level>",
-    )
     config.output.write_data(pl.concat(dfs))
 
 

anonymize/models/sources.py

@@ -1,3 +1,4 @@
+import sys
 from typing import Literal, Union
 from pydantic import BaseModel, Field
 
@@ -62,6 +63,12 @@ def __iter__(self):
     def __next__(self):
         data = self.read_data()
         if data.first().collect().is_empty():
+            # Bring back the logger
+            logger.add(
+                sys.stdout,
+                colorize=True,
+                format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <level>{message}</level>",
+            )
             raise StopIteration
         return data