From 2066e20bbc50ad811c85a20e61e62b8b98258189 Mon Sep 17 00:00:00 2001
From: Adam <adam.remaki@aphp.fr>
Date: Thu, 28 Sep 2023 17:41:52 +0200
Subject: [PATCH] Quick fix

---
 edsteva/probes/utils/prepare_df.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/edsteva/probes/utils/prepare_df.py b/edsteva/probes/utils/prepare_df.py
index 0b42634..e3a28d0 100644
--- a/edsteva/probes/utils/prepare_df.py
+++ b/edsteva/probes/utils/prepare_df.py
@@ -56,25 +56,33 @@ def prepare_visit_occurrence(
         invalid_naming="supprimé",
     )
 
-    if stay_sources:
+    # 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",
         )
 
-    if provenance_sources:
+    # 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",
         )
 
-    if length_of_stays:
+    if length_of_stays and isinstance(length_of_stays, list):
         visit_occurrence = filter_table_by_length_of_stay(
             visit_occurrence=visit_occurrence, length_of_stays=length_of_stays
         )
@@ -103,7 +111,7 @@ def prepare_visit_occurrence(
             target_col="stay_type",
         )
 
-    if age_ranges:
+    if age_ranges and isinstance(age_ranges, list):
         visit_occurrence = visit_occurrence.merge(person, on="person_id")
         visit_occurrence = filter_table_by_age(
             visit_occurrence=visit_occurrence,
@@ -184,7 +192,7 @@ def prepare_measurement(
             )
         measurement = get_framework(measurement).concat(measurement_by_terminology)
 
-    if concepts_sets:
+    if concepts_sets and isinstance(concepts_sets, (dict, str)):
         measurement_by_terminology = []
         for standard_terminology in standard_terminologies:
             measurement_by_terminology.append(