From e8f84796d9ed20bb3da87e92844a5760f166c9fe Mon Sep 17 00:00:00 2001
From: Adam <adam.remaki@aphp.fr>
Date: Thu, 28 Sep 2023 11:39:32 +0200
Subject: [PATCH] fix: Rename the columns if the variable is a bool

---
 edsteva/probes/utils/prepare_df.py | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/edsteva/probes/utils/prepare_df.py b/edsteva/probes/utils/prepare_df.py
index 940a8fc..cb1d4ec 100644
--- a/edsteva/probes/utils/prepare_df.py
+++ b/edsteva/probes/utils/prepare_df.py
@@ -56,21 +56,29 @@ def prepare_visit_occurrence(
         invalid_naming="supprimé",
     )
 
+    # Filter stay source (MCO, PSY, SSR)
+    visit_occurrence = visit_occurrence.rename(
+        columns={"stay_source_value": "stay_source"}
+    )
     if stay_sources and isinstance(stay_sources, (dict, str)):
         visit_occurrence = filter_table_by_type(
             table=visit_occurrence,
             table_name="visit_occurrence",
             type_groups=stay_sources,
-            source_col="stay_source_value",
+            source_col="stay_source",
             target_col="stay_source",
         )
 
+    # Filter provenance source
+    visit_occurrence = visit_occurrence.rename(
+        columns={"provenance_source_value": "provenance_source"}
+    )
     if provenance_sources and isinstance(provenance_sources, (dict, str)):
         visit_occurrence = filter_table_by_type(
             table=visit_occurrence,
             table_name="visit_occurrence",
             type_groups=provenance_sources,
-            source_col="provenance_source_value",
+            source_col="provenance_source",
             target_col="provenance_source",
         )
 
@@ -79,10 +87,8 @@ def prepare_visit_occurrence(
             visit_occurrence=visit_occurrence, length_of_stays=length_of_stays
         )
 
-    visit_occurrence = visit_occurrence.rename(
-        columns={"visit_source_value": "stay_type", "visit_start_datetime": "date"}
-    )
-
+    # Filter date
+    visit_occurrence = visit_occurrence.rename(columns={"visit_start_datetime": "date"})
     visit_occurrence = filter_table_by_date(
         table=visit_occurrence,
         table_name="visit_occurrence",
@@ -94,6 +100,10 @@ def prepare_visit_occurrence(
         cost = cost[["visit_occurrence_id", "drg_source"]].drop_duplicates()
         visit_occurrence = visit_occurrence.merge(cost, on="visit_occurrence_id")
 
+    # Filter stay type (hospitalisation, consultation, urgence)
+    visit_occurrence = visit_occurrence.rename(
+        columns={"visit_source_value": "stay_type"}
+    )
     if stay_types and isinstance(stay_types, (dict, str)):
         visit_occurrence = filter_table_by_type(
             table=visit_occurrence,
@@ -400,12 +410,11 @@ def prepare_note(
             "note_text",
         ]
     ]
-    note = note.rename(
-        columns={"note_class_source_value": "note_type", "note_datetime": "date"}
-    )
     note = note[~(note["note_text"].isna())]
     note = note.drop(columns=["note_text"])
 
+    # Filter date
+    note = note.rename(columns={"note_datetime": "date"})
     note = filter_table_by_date(
         table=note,
         table_name="note",
@@ -417,7 +426,8 @@ def prepare_note(
         table=note, table_name="note", valid_naming="Actif"
     )
 
-    # Add note type
+    # Filter note type
+    note = note.rename(columns={"note_class_source_value": "note_type"})
     if note_types and isinstance(note_types, (dict, str)):
         note = filter_table_by_type(
             table=note,