Merge pull request #18 from TidierOrg/add-cat_expand

bump lts, add cat_expand

Project.toml

@@ -1,7 +1,7 @@
 name = "TidierCats"
 uuid = "79ddc9fe-4dbf-4a56-a832-df41fb326d23"
 authors = ["Daniel Rizk"]
-version = "0.1.2"
+version = "0.2.0"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
@@ -13,7 +13,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 CategoricalArrays = "0.10, 1.0"
 DataFrames = "1.5"
 Reexport = "0.2, 1"
-julia = "1.9"
+julia = "1.10"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/TidierCats.jl

@@ -8,7 +8,8 @@ using Reexport
 @reexport using CategoricalArrays
 
 export cat_rev, cat_relevel, cat_infreq, cat_lump, cat_reorder, cat_collapse, cat_lump_min, cat_lump_prop
-export as_categorical, as_integer, cat_replace_missing, cat_other, cat_recode
+export as_categorical, as_integer, cat_replace_missing, cat_other, cat_recode, cat_expand
+export cat_ages
 include("catsdocstrings.jl")
 
 """
@@ -130,9 +131,17 @@ end
 $docstring_as_categorical
 """
 function as_categorical(arr::AbstractArray)
-  return CategoricalArray(map(x -> ismissing(x) ? missing : x, arr)) 
+    T = eltype(arr)
+    if T <: Number
+        # keep numeric data as numeric categories
+        return CategoricalArray{Union{Missing, T}}(arr)
+    else
+        # fallback: treat them as strings
+        return CategoricalArray(map(x -> ismissing(x) ? missing : string(x), arr))
+    end
 end
 
+
 """
 $docstring_cat_reorder
 """
@@ -329,5 +338,33 @@ function cat_recode(f::Union{CategoricalArray, AbstractVector}; kwargs...)
 end
 
 
+"""
+$docstring_cat_expand
+"""
+function cat_expand(f::CategoricalArray, new_levels...; after=Inf)
+    # Get the current levels of the factor
+    current_levels = levels(f)
+
+    # Filter out new levels that already exist
+    unique_new_levels = [level for level in new_levels if level ∉ current_levels]
+
+    # Decide where to place the new levels based on the `after` argument
+    if after == Inf
+        # Append the new levels at the end if `after` is Inf (default)
+        expanded_levels = vcat(current_levels, unique_new_levels)
+    elseif after == 0
+        # Prepend the new levels at the beginning if `after` is 0
+        expanded_levels = vcat(unique_new_levels, current_levels)
+    else
+        # Insert the new levels after the specified index
+        expanded_levels = vcat(current_levels[1:after], unique_new_levels, current_levels[after+1:end])
+    end
+
+    # Update the levels of the categorical array
+    levels!(f, expanded_levels)
+
+    return f
+end
+
 
 end

src/catsdocstrings.jl

@@ -408,3 +408,31 @@ julia> cat_other(cat_array, drop = ["A", "B"])
  "E"
 ```
 """
+
+const docstring_cat_expand = 
+"""
+    cat_expand(cat_array::CategoricalArray, new_levels...; after=Inf)
+
+Expands the levels in a categorical array by adding new levels at a specified position.
+
+# Arguments
+- `cat_array`: Categorical array to expand levels
+- `new_levels`: New levels to be added to the categorical array
+- `after`: Position after which to insert the new levels. Default is Inf, which appends the new levels at the end.
+
+# Returns
+Categorical array with the new levels added.
+
+# Examples
+```julia
+julia> cats = CategoricalArray(["a", "b", "c", "a", "c", "b"]);
+
+julia> println("Original levels: ", levels(cats))
+Original levels: ["a", "b", "c"]
+
+julia> cats = cat_expand(f, "d", "e", "f");
+
+julia> println("Expanded levels: ", levels(cats))
+Expanded levels: ["a", "b", "c", "d", "e", "f"]
+```
+"""