Quick fix

aphp · Sep 28, 2023 · 2066e20 · 2066e20
1 parent bafdac3
commit 2066e20
Showing 1 changed file with 15 additions and 7 deletions.
diff --git 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(