diff --git a/src/utils.jl b/src/utils.jl index bab266c..e8ff7dc 100644 --- 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 index 2843b01..077c2c7 100644 --- 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