From c9b1c5f30cbcfabe8d5147d1b63ed5c8734a5c20 Mon Sep 17 00:00:00 2001
From: Elias Carvalho <73039601+eliascarv@users.noreply.github.com>
Date: Fri, 2 Feb 2024 15:20:33 -0300
Subject: [PATCH] Rename: Add support for modification function (#267)

* Rename: Add support for modification function

* Update compat
---
 src/transforms/rename.jl  | 31 ++++++++++++++++++++++---------
 test/Project.toml         |  3 ++-
 test/transforms/rename.jl |  6 ++++++
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/transforms/rename.jl b/src/transforms/rename.jl
index 5396833c..fc08fbc3 100644
--- a/src/transforms/rename.jl
+++ b/src/transforms/rename.jl
@@ -6,8 +6,12 @@
     Rename(:col₁ => :newcol₁, :col₂ => :newcol₂, ..., :colₙ => :newcolₙ)
     Rename([:col₁ => :newcol₁, :col₂ => :newcol₂, ..., :colₙ => :newcolₙ])
 
-The transform that renames `col₁`, `col₂`, ..., `colₙ`
-to `newcol₁`, `newcol₂`, ..., `newcolₙ`.
+Renames the columns `col₁`, `col₂`, ..., `colₙ` to `newcol₁`, `newcol₂`, ..., `newcolₙ`.
+
+    Rename(fun)
+
+Renames the table columns using the modification function `fun` that takes a 
+string as input and returns another string with the new name.
 
 # Examples
 
@@ -18,17 +22,22 @@ Rename("a" => "x", "c" => "y")
 Rename([1 => "x", 3 => "y"])
 Rename([:a => "x", :c => "y"])
 Rename(["a", "c"] .=> [:x, :y])
+Rename(nm -> nm * "_suffix")
 ```
 """
-struct Rename{S<:ColumnSelector} <: StatelessFeatureTransform
+struct Rename{S<:ColumnSelector,N} <: StatelessFeatureTransform
   selector::S
-  newnames::Vector{Symbol}
-  function Rename(selector::S, newnames) where {S<:ColumnSelector}
-    _assert(allunique(newnames), "new names must be unique")
-    new{S}(selector, newnames)
+  newnames::N
+  function Rename(selector::S, newnames::N) where {S<:ColumnSelector,N}
+    if newnames isa AbstractVector
+      _assert(allunique(newnames), "new names must be unique")
+    end
+    new{S,N}(selector, newnames)
   end
 end
 
+Rename(fun) = Rename(AllSelector(), fun)
+
 Rename(pairs::Pair{C,Symbol}...) where {C<:Column} = Rename(selector(first.(pairs)), collect(last.(pairs)))
 
 Rename(pairs::Pair{C,S}...) where {C<:Column,S<:AbstractString} =
@@ -41,13 +50,17 @@ Rename(pairs::AbstractVector{Pair{C,S}}) where {C<:Column,S<:AbstractString} =
 
 isrevertible(::Type{<:Rename}) = true
 
+_newnames(newnames::AbstractVector{Symbol}, snames) = newnames
+_newnames(fun, snames) = [Symbol(fun(string(name))) for name in snames]
+
 function applyfeat(transform::Rename, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols)
   snames = transform.selector(names)
-  _assert(transform.newnames ⊈ setdiff(names, snames), "duplicate names")
+  tnames = _newnames(transform.newnames, snames)
+  _assert(tnames ⊈ setdiff(names, snames), "duplicate names")
 
-  mapnames = Dict(zip(snames, transform.newnames))
+  mapnames = Dict(zip(snames, tnames))
   newnames = [get(mapnames, nm, nm) for nm in names]
   columns = [Tables.getcolumn(cols, nm) for nm in names]
 
diff --git a/test/Project.toml b/test/Project.toml
index 1cc1b33c..cee8ce0b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -19,4 +19,5 @@ TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
-PairPlots = "2.1"
+CairoMakie = "=0.11.6"
+PairPlots = "=2.3.1"
diff --git a/test/transforms/rename.jl b/test/transforms/rename.jl
index cf246171..0489f95f 100644
--- a/test/transforms/rename.jl
+++ b/test/transforms/rename.jl
@@ -146,6 +146,12 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
+  T = Rename(nm -> nm * "_test")
+  n, c = apply(T, t)
+  @test Tables.schema(n).names == (:a_test, :b_test, :c_test, :d_test)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
   # throws
   @test_throws AssertionError Rename(:a => :x, :b => :x)
   @test_throws AssertionError apply(Rename(:a => :c, :b => :d), t)