From 02d391adae1e417ac6d99f92da54d9054123bfd0 Mon Sep 17 00:00:00 2001
From: enamya <aymaneboumaaza@gmail.com>
Date: Thu, 11 Apr 2024 23:32:58 +0000
Subject: [PATCH] Fix double logging issue

---
 anonymize/__main__.py       | 9 ---------
 anonymize/models/sources.py | 7 +++++++
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/anonymize/__main__.py b/anonymize/__main__.py
index da5fdb1..760672a 100644
--- a/anonymize/__main__.py
+++ b/anonymize/__main__.py
@@ -1,4 +1,3 @@
-import sys
 from anonymize import load_config, Config
 
 from loguru import logger
@@ -9,9 +8,6 @@
 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.")
@@ -19,11 +15,6 @@ def main(config: Config):
             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))
 
 
diff --git a/anonymize/models/sources.py b/anonymize/models/sources.py
index 0daa257..dfa9348 100644
--- a/anonymize/models/sources.py
+++ b/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