Skip to content

Commit

Permalink
Add RegroupedComponentSelector to regroup ListComponentSelector
Browse files Browse the repository at this point in the history
  • Loading branch information
GabrielKS committed Oct 30, 2024
1 parent 62c0de4 commit ed5c391
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 1 deletion.
47 changes: 46 additions & 1 deletion src/component_selector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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...)
29 changes: 29 additions & 0 deletions test/test_component_selector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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

0 comments on commit ed5c391

Please sign in to comment.