From ed5c3919aaea0405b3ad37cfe0e79c5d5cec583e Mon Sep 17 00:00:00 2001
From: GabrielKS <23368820+GabrielKS@users.noreply.github.com>
Date: Tue, 29 Oct 2024 23:15:44 -0600
Subject: [PATCH] Add `RegroupedComponentSelector` to regroup
 `ListComponentSelector`

---
 src/component_selector.jl       | 47 ++++++++++++++++++++++++++++++++-
 test/test_component_selector.jl | 29 ++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/src/component_selector.jl b/src/component_selector.jl
index eb0737dd..ac6c3fa1 100644
--- a/src/component_selector.jl
+++ b/src/component_selector.jl
@@ -299,6 +299,33 @@ function get_components(
     return FlattenIteratorWrapper(my_supertype, sub_components)
 end
 
+# Rebuilding
+"""
+Returns a `ComponentSelector` functionally identical to the input `selector` except with the
+changes to its fields specified in the keyword arguments. For `ListComponentSelector`, if a
+`groupby` option is specified, the return type will be a `RegroupedComponentSelector`
+instead of a `ListComponentSelector`.
+
+# Examples
+Suppose you have a selector with manual groups and you want to group by `:each`:
+```julia
+sel = make_selector(make_selector(ThermalStandard), make_selector(RenewableDispatch))
+sel_each = rebuild_selector(sel; groupby = :each)  # will be a RegroupedComponentSelector
+```
+"""
+function rebuild_selector(selector::ListComponentSelector;
+    name = nothing, groupby = nothing)
+    # Handle the easy stuff first
+    selector_data =
+        Dict(key => getfield(selector, key) for key in fieldnames(typeof(selector)))
+    isnothing(name) || (selector_data[:name] = name)
+    rebuilt = ListComponentSelector(; selector_data...)
+
+    # Wrap in a RegroupedComponentSelector if we need to
+    isnothing(groupby) && return rebuilt
+    return RegroupedComponentSelector(rebuilt, groupby)
+end
+
 # TypeComponentSelector
 "`PluralComponentSelector` represented by a type of component."
 @kwdef struct TypeComponentSelector <: DynamicallyGroupedComponentSelector
@@ -371,7 +398,6 @@ FilterComponentSelector(
         string(filter_func) * COMPONENT_NAME_DELIMITER * subtype_to_string(component_type),
     )
 
-# Signature 2: put the filter function first for consistency with non-`ComponentSelector` `get_components`
 """
 Make a ComponentSelector from a filter function and a type of component. The filter function
 must accept instances of `component_type` as a sole argument and return a `Bool`. Optionally
@@ -401,3 +427,22 @@ function get_components(
     components = get_components(combo_filter, selector.component_type, sys)
     return components
 end
+
+# RegroupedComponentSelector
+"`PluralComponentSelector` that wraps another `ComponentSelector` and applies dynamic grouping."
+@kwdef struct RegroupedComponentSelector <: DynamicallyGroupedComponentSelector
+    wrapped_selector::ComponentSelector
+    groupby::Union{Symbol, Function}
+
+    RegroupedComponentSelector(
+        wrapped_selector::ComponentSelector,
+        groupby::Union{Symbol, Function},
+    ) = new(wrapped_selector, validate_groupby(groupby))
+end
+
+# Naming
+get_name(selector::RegroupedComponentSelector) = get_name(selector.wrapped_selector)
+
+# Contents
+get_components(selector::RegroupedComponentSelector, sys::SystemLike; kwargs...) =
+    get_components(selector.wrapped_selector, sys, kwargs...)
diff --git a/test/test_component_selector.jl b/test/test_component_selector.jl
index 8137a706..564f1e05 100644
--- a/test/test_component_selector.jl
+++ b/test/test_component_selector.jl
@@ -259,6 +259,25 @@ end
     end
 end
 
+@testset "Test RegroupedComponentSelector" begin
+    @testset for test_sys in [cstest_make_components(), cstest_make_system_data()]
+        comp_ent_1 = IS.make_selector(IS.TestComponent, "Component1")
+        comp_ent_2 = IS.make_selector(IS.AdditionalTestComponent, "Component3")
+        test_list_ent = IS.ListComponentSelector((comp_ent_1, comp_ent_2), nothing)
+        test_sel = IS.RegroupedComponentSelector(test_list_ent, :all)
+
+        # Equality
+        @test IS.RegroupedComponentSelector(test_list_ent, :all) == test_sel
+
+        # Naming
+        @test IS.get_name(test_sel) == IS.get_name(test_list_ent)
+
+        # Contents
+        @test Set(collect(get_components_rt(test_sel, test_sys))) ==
+              Set(collect(get_components_rt(test_list_ent, test_sys)))
+    end
+end
+
 @testset "Test DynamicallyGroupedComponentSelector grouping" begin
     # We'll use TypeComponentSelector as the token example
     @assert IS.TypeComponentSelector <: IS.DynamicallyGroupedComponentSelector
@@ -301,6 +320,7 @@ end
     sel1::IS.NameComponentSelector =
         IS.make_selector(IS.TestComponent, "Component1"; name = "oldname")
     sel2::IS.TypeComponentSelector = IS.make_selector(IS.TestComponent; groupby = :all)
+    sel3::IS.ListComponentSelector = IS.make_selector(sel1, sel2; name = "oldname")
 
     @test IS.rebuild_selector(sel1; name = "newname") ==
           IS.make_selector(IS.TestComponent, "Component1"; name = "newname")
@@ -310,4 +330,13 @@ end
           IS.make_selector(IS.TestComponent; name = "newname", groupby = :all)
     @test IS.rebuild_selector(sel2; name = "newname", groupby = :each) ==
           IS.make_selector(IS.TestComponent; name = "newname", groupby = :each)
+
+    @test IS.rebuild_selector(sel3; name = "newname") ==
+          IS.make_selector(sel1, sel2; name = "newname")
+    regrouped = IS.rebuild_selector(sel3; name = "newname", groupby = :all)
+    for test_sys in [cstest_make_components(), cstest_make_system_data()]
+        @test Set(collect(get_components_rt(regrouped, test_sys))) ==
+              Set(collect(get_components_rt(sel3, test_sys)))
+        @test length(IS.get_groups(regrouped, test_sys)) == 1
+    end
 end