better handling of missing in wrapped type

TARGENE · Mar 29, 2024 · fa165e0 · fa165e0
1 parent 1869dbb
commit fa165e0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/utils.jl b/src/utils.jl
@@ -37,7 +37,7 @@ end
 
 wrapped_type(x) = x
 wrapped_type(x::Type{<:CategoricalValue{T,}}) where T = T
-
+wrapped_type(x::Type{Union{Missing, T}}) where T = wrapped_type(T)
 """
 Uses the values found in the dataset to create a new estimand with adjusted values.
 """
@@ -137,7 +137,7 @@ end
 
 function make_categorical!(dataset, colnames; infer_ordered=false)
     for colname in colnames
-        make_categorical!(dataset, colname;infer_ordered=infer_ordered)
+        make_categorical!(dataset, colname; infer_ordered=infer_ordered)
     end
 end
 

diff --git a/test/utils.jl b/test/utils.jl
@@ -141,6 +141,16 @@ end
     data.W₁ = [1, 2, missing, 4, 5]
     sample_ids = TargetedEstimation.sample_ids_from_variables(data, variables)
     @test sample_ids == [2]
+    # wrapped_ype
+    col = categorical(["AC", "CC"])
+    @test TargetedEstimation.wrapped_type(eltype(col)) == String
+    col = categorical(["AC", "CC", missing])
+    @test TargetedEstimation.wrapped_type(eltype(col)) == String
+    col = [1, missing, 0.3]
+    @test TargetedEstimation.wrapped_type(eltype(col)) == Float64
+    col = [1, 2, 3]
+    @test TargetedEstimation.wrapped_type(eltype(col)) == Int64
+
 end
 
 @testset "Test make_categorical! and make_float!" begin