Merge pull request #489 from luke-kiernan/issue-485

get_supplemental_attribute methods with component type input

Author: jd-lara

src/component.jl

@@ -87,8 +87,8 @@ function _get_supplemental_attributes(
         get_supplemental_attribute(mgr, x) for
         x in list_associated_supplemental_attribute_uuids(
             mgr.associations,
-            component;
-            attribute_type = supplemental_attribute_type,
+            component,
+            supplemental_attribute_type,
         )
     ]
 end
@@ -118,8 +118,8 @@ function _get_supplemental_attributes(
     attrs = Vector{supplemental_attribute_type}()
     for uuid in list_associated_supplemental_attribute_uuids(
         mgr.associations,
-        component;
-        attribute_type = supplemental_attribute_type,
+        component,
+        supplemental_attribute_type,
     )
         attribute = get_supplemental_attribute(mgr, uuid)
         if filter_func(attribute)

src/supplemental_attribute_associations.jl

@@ -288,6 +288,7 @@ function has_association(
     )
 end
 
+# component UUIDS from attribute
 const _QUERY_LIST_ASSOCIATED_COMP_UUIDS = """
     SELECT DISTINCT component_uuid
     FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
@@ -309,6 +310,38 @@ function list_associated_component_uuids(
     return Base.UUID.(table.component_uuid)
 end
 
+# attribute UUIDS from component
+const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS = """
+    SELECT DISTINCT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE component_uuid = ?
+"""
+
+"""
+Return the distinct attribute UUIDs associated with the component.
+"""
+function list_associated_supplemental_attribute_uuids(
+    associations::SupplementalAttributeAssociations,
+    component::InfrastructureSystemsComponent,
+    ::Nothing,
+)
+    params = (string(get_uuid(component)),)
+    table = Tables.columntable(
+        _execute_cached(
+            associations,
+            _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS,
+            params,
+        ),
+    )
+    return Base.UUID.(table.attribute_uuid)
+end
+
+list_associated_supplemental_attribute_uuids(
+    associations::SupplementalAttributeAssociations,
+    component::InfrastructureSystemsComponent,
+) = list_associated_supplemental_attribute_uuids(associations, component, nothing)
+
+# component UUIDS from attribute plus component type
 const _QUERY_LIST_ASSOCIATED_COMP_UUIDS_BY_C_TYPE = StringTemplates.@template """
     SELECT DISTINCT component_uuid
     FROM $table_name
@@ -336,6 +369,36 @@ function list_associated_component_uuids(
     return Base.UUID.(table.component_uuid)
 end
 
+# attribute UUIDS from component plus attribute type
+const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_A_TYPE =
+    StringTemplates.@template """
+    SELECT DISTINCT attribute_uuid
+    FROM $table_name
+    WHERE component_uuid = ? AND $attribute_type_clause
+"""
+
+"""
+Return the distinct attribute UUIDs associated with the component, filter by attribute type.
+"""
+function list_associated_supplemental_attribute_uuids(
+    associations::SupplementalAttributeAssociations,
+    component::InfrastructureSystemsComponent,
+    attribute_type::Type{<:SupplementalAttribute},
+)
+    params = [string(get_uuid(component))]
+    attribute_type_clause = _get_type_clause!(params, attribute_type, "attribute_type")
+    query = StringTemplates.render(
+        _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_A_TYPE;
+        table_name = SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
+        attribute_type_clause = attribute_type_clause,
+    )
+    table = Tables.columntable(
+        _execute_cached(associations, query, params),
+    )
+    return Base.UUID.(table.attribute_uuid)
+end
+
+# component UUIDS from attribute type driver function
 """
 Return the distinct component UUIDs associated with the attribute type.
 """
@@ -352,6 +415,27 @@ function list_associated_component_uuids(
     return _list_associated_component_uuids(associations, subtypes)
 end
 
+# attribute UUIDS from component type driver function
+"""
+Return the distinct attribute UUIDs associated with the component type.
+"""
+function list_associated_supplemental_attribute_uuids(
+    associations::SupplementalAttributeAssociations,
+    component_type::Type{<:InfrastructureSystemsComponent},
+    ::Nothing,
+)
+    if isconcretetype(component_type)
+        return _list_associated_supplemental_attribute_uuids(
+            associations,
+            (component_type,),
+        )
+    end
+
+    subtypes = get_all_concrete_subtypes(component_type)
+    return _list_associated_supplemental_attribute_uuids(associations, subtypes)
+end
+
+# component UUIDS from attribute type plus component type
 const _QUERY_LIST_ASSOCIATED_COMP_UUIDS_BY_TYPES = StringTemplates.@template """
     SELECT DISTINCT component_uuid
     FROM $table_name
@@ -380,6 +464,36 @@ function list_associated_component_uuids(
     return Base.UUID.(table.component_uuid)
 end
 
+# attribute UUIDS from attribute type plus component type
+const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_TYPES =
+    StringTemplates.@template """
+    SELECT DISTINCT attribute_uuid
+    FROM $table_name
+    WHERE $attribute_type_clause AND $component_type_clause
+"""
+
+"""
+Return the distinct attribute UUIDs associated with the component type, filter by
+attribute_type.
+"""
+function list_associated_supplemental_attribute_uuids(
+    associations::SupplementalAttributeAssociations,
+    component_type::Type{<:InfrastructureSystemsComponent},
+    attribute_type::Type{<:SupplementalAttribute},
+)
+    params = String[]
+    attribute_type_clause = _get_type_clause!(params, attribute_type, "attribute_type")
+    component_type_clause = _get_type_clause!(params, component_type, "component_type")
+    query = StringTemplates.render(
+        _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_TYPES;
+        table_name = SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
+        attribute_type_clause = attribute_type_clause,
+        component_type_clause = component_type_clause,
+    )
+    table = Tables.columntable(_execute_cached(associations, query, params))
+    return Base.UUID.(table.attribute_uuid)
+end
+
 const _QUERY_LIST_ASSOCIATED_PAIR_UUIDS = StringTemplates.@template """
     SELECT DISTINCT component_uuid, attribute_uuid
     FROM $table_name
@@ -431,6 +545,7 @@ function _get_type_clause!(
     return type_clause
 end
 
+# component UUIDs from attribute type implementation
 const _QUERY_LIST_ASSOCIATED_COMP_UUIDS_BY_ONE_TYPE = """
     SELECT DISTINCT component_uuid
     FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
@@ -468,32 +583,41 @@ function _list_associated_component_uuids(
     return Base.UUID.(table.component_uuid)
 end
 
-const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS = StringTemplates.@template """
-    SELECT attribute_uuid
-    FROM $table_name
-    WHERE $where_clause
-"""
+# attribute UUIDs from component type implementation
+const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_ONE_TYPE = """
+    SELECT DISTINCT attribute_uuid
+    FROM $SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME
+    WHERE component_type = ?
 """
-Return the supplemental attribute UUIDs associated with the component and attribute type.
+
+const _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_MULTIPLE_TYPES =
+    StringTemplates.@template """
+    SELECT DISTINCT attribute_uuid
+    FROM $table_name
+    WHERE component_type IN ($placeholder)
 """
-function list_associated_supplemental_attribute_uuids(
+
+function _list_associated_supplemental_attribute_uuids(
     associations::SupplementalAttributeAssociations,
-    component::InfrastructureSystemsComponent;
-    attribute_type::Union{Nothing, Type{<:SupplementalAttribute}} = nothing,
+    component_types,
 )
-    c_str = "component_uuid = ?"
-    params = [string(get_uuid(component))]
-    if isnothing(attribute_type)
-        where_clause = c_str
+    len = length(component_types)
+    if len == 0
+        # This would require an abstract type with no subtypes. Just here for completeness.
+        return Base.UUID[]
+    elseif len == 1
+        query = _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_ONE_TYPE
+        params = (string(nameof(first(component_types))),)
     else
-        a_str = _get_attribute_type_string!(params, attribute_type)
-        where_clause = "$c_str AND $a_str"
+        placeholder = chop(repeat("?,", length(component_types)))
+        params = Tuple(string(nameof(type)) for type in component_types)
+        query = StringTemplates.render(
+            _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS_BY_MULTIPLE_TYPES;
+            table_name = SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
+            placeholder = placeholder,
+        )
     end
-    query = StringTemplates.render(
-        _QUERY_LIST_ASSOCIATED_SUPPLEMENTAL_ATTRIBUTE_UUIDS;
-        table_name = SUPPLEMENTAL_ATTRIBUTE_TABLE_NAME,
-        where_clause = where_clause,
-    )
+
     table = Tables.columntable(_execute_cached(associations, query, params))
     return Base.UUID.(table.attribute_uuid)
 end

src/supplemental_attribute_manager.jl

@@ -248,6 +248,18 @@ function list_associated_component_uuids(
     return list_associated_component_uuids(mgr.associations, attribute_type, component_type)
 end
 
+function list_associated_supplemental_attribute_uuids(
+    mgr::SupplementalAttributeManager,
+    component_type::Type{<:InfrastructureSystemsComponent},
+    attribute_type::Union{Nothing, Type{<:SupplementalAttribute}},
+)
+    return list_associated_supplemental_attribute_uuids(
+        mgr.associations,
+        component_type,
+        attribute_type,
+    )
+end
+
 function serialize(mgr::SupplementalAttributeManager)
     return Dict(
         "associations" => to_records(mgr.associations),

src/system_data.jl

@@ -1087,6 +1087,31 @@ function get_associated_components(
     ]
 end
 
+"""
+Return all supplemental attributes associated with the components of the given type, optionally filtered by `attribute_type`.
+
+# Arguments
+- `data::SystemData`: the `SystemData` to search
+- `component_type`::Type{<:InfrastructureSystemsComponent}: Only return attributes
+ associated with the components of this type.
+- `attribute_type`::Union{Nothing, Type{<:SupplementalAttribute}}`: Optional, type of the
+  attributes to return. Can be concrete or abstract.
+"""
+function get_associated_supplemental_attributes(
+    data::SystemData,
+    component_type::Type{<:InfrastructureSystemsComponent};
+    attribute_type::Union{Nothing, Type{<:SupplementalAttribute}} = nothing,
+)
+    return [
+        get_supplemental_attribute(data, x) for x in
+        list_associated_supplemental_attribute_uuids(
+            data.supplemental_attribute_manager,
+            component_type,
+            attribute_type,
+        )
+    ]
+end
+
 """
 Return all components associated with the attribute that match `component_type`.
 

test/test_function_data.jl

@@ -209,10 +209,11 @@ end
 end
 
 @testset "Test FunctionData printing" begin
+    type_prefix = @static VERSION >= v"1.12" ? "" : "@NamedTuple{x::Float64, y::Float64}"
     repr_answers = [
         "InfrastructureSystems.LinearFunctionData(5.0, 1.0)",
         "InfrastructureSystems.QuadraticFunctionData(2.0, 3.0, 4.0)",
-        "InfrastructureSystems.PiecewiseLinearData(@NamedTuple{x::Float64, y::Float64}[(x = 1.0, y = 1.0), (x = 3.0, y = 5.0), (x = 5.0, y = 10.0)])",
+        "InfrastructureSystems.PiecewiseLinearData($(type_prefix)[(x = 1.0, y = 1.0), (x = 3.0, y = 5.0), (x = 5.0, y = 10.0)])",
         "InfrastructureSystems.PiecewiseStepData([1.0, 3.0, 5.0], [2.0, 2.5])",
     ]
     plain_answers = [

test/test_random_seed.jl

@@ -1,4 +1,3 @@
-IS = InfrastructureSystems
 @testset "Random Seed" begin
     @test isa(IS.get_random_seed(), Int)
 end

test/test_system_data.jl

@@ -396,6 +396,48 @@ end
     @test length(attributes) == 0
 end
 
+@testset "Test get supplemental attribute from component type" begin
+    data = IS.SystemData()
+
+    attr1 = IS.TestSupplemental(; value = 1.0)
+    attr2 = IS.TestSupplemental(; value = 2.0)
+    geo_attr1 = IS.GeographicInfo()
+    geo_attr2 = IS.GeographicInfo(; geo_json = Dict{String, Any}("foo" => 5))
+    components_to_attributes = Dict(
+        (IS.TestComponent, "component1") => [geo_attr1, attr1],
+        (IS.TestComponent, "component2") => [attr1],
+        (IS.AdditionalTestComponent, "component3") => [geo_attr1, attr2, geo_attr2],
+    )
+
+    for (comp_info, attributes) in components_to_attributes
+        component_type, name = comp_info
+        component = component_type(name, 5)
+        IS.add_component!(data, component)
+        for attribute in attributes
+            IS.add_supplemental_attribute!(data, component, attribute)
+        end
+    end
+
+    # no duplicate attr1, despite occuring twice.
+    test_comp_attrs = IS.get_associated_supplemental_attributes(data, IS.TestComponent)
+    @test length(test_comp_attrs) == 2 && attr1 in test_comp_attrs &&
+          geo_attr1 in test_comp_attrs
+
+    add_test_comp_attrs =
+        IS.get_associated_supplemental_attributes(data, IS.AdditionalTestComponent)
+    @test length(add_test_comp_attrs) == 3 && attr2 in add_test_comp_attrs &&
+          geo_attr1 in add_test_comp_attrs && geo_attr2 in add_test_comp_attrs
+
+    # test additional attribute_type kwarg.
+    add_test_comp_geo_attrs = IS.get_associated_supplemental_attributes(
+        data,
+        IS.AdditionalTestComponent;
+        attribute_type = IS.GeographicInfo,
+    )
+    @test length(add_test_comp_geo_attrs) == 2 && geo_attr1 in add_test_comp_geo_attrs &&
+          geo_attr2 in add_test_comp_geo_attrs
+end
+
 @testset "Test retrieval of supplemental_attributes" begin
     data = IS.SystemData()
     geo_supplemental_attribute = IS.GeographicInfo()