Add get_supplemental_attributes for components

NREL-Sienna · Jan 8, 2024 · 3bdfd53 · 3bdfd53
1 parent f0db36a
commit 3bdfd53
Show file tree

Hide file tree

Showing 12 changed files with 243 additions and 117 deletions.
diff --git a/src/InfrastructureSystems.jl b/src/InfrastructureSystems.jl
@@ -75,7 +75,7 @@ Subtypes may contain time series. Which requires
 All subtypes must include an instance of ComponentUUIDs in order to track
 components attached to each attribute.
 """
-abstract type InfrastructureSystemsSupplementalAttribute <: InfrastructureSystemsType end
+abstract type SupplementalAttribute <: InfrastructureSystemsType end
 
 """
 Return the internal time_series storage container or nothing, if the type doesn't store
@@ -119,9 +119,10 @@ include("time_series_parser.jl")
 include("containers.jl")
 include("component_uuids.jl")
 include("supplemental_attribute.jl")
-include("supplemental_attributes.jl")
 include("supplemental_attributes_container.jl")
+include("supplemental_attributes.jl")
 include("components.jl")
+include("iterators.jl")
 include("geographic_supplemental_attribute.jl")
 include("generated/includes.jl")
 include("single_time_series.jl")

diff --git a/src/component.jl b/src/component.jl
@@ -301,7 +301,7 @@ Attach an attribute to a component.
 function attach_supplemental_attribute!(
     component::InfrastructureSystemsComponent,
     attribute::T,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     attribute_container = get_supplemental_attributes_container(component)
 
     if !haskey(attribute_container, T)
@@ -328,7 +328,7 @@ Return true if the component has supplemental attributes of the given type.
 function has_supplemental_attributes(
     ::Type{T},
     component::InfrastructureSystemsComponent,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     supplemental_attributes = get_supplemental_attributes_container(component)
     if !isconcretetype(T)
         for (k, v) in supplemental_attributes
@@ -366,7 +366,7 @@ end
 function detach_supplemental_attribute!(
     component::InfrastructureSystemsComponent,
     attribute::T,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     container = get_supplemental_attributes_container(component)
     if !haskey(container, T)
         throw(
@@ -381,3 +381,44 @@ function detach_supplemental_attribute!(
     end
     return
 end
+
+"""
+Returns an iterator of supplemental_attributes. T can be concrete or abstract.
+Call collect on the result if an array is desired.
+
+# Arguments
+
+  - `T`: supplemental_attribute type
+  - `supplemental_attributes::SupplementalAttributes`: SupplementalAttributes in the system
+  - `filter_func::Union{Nothing, Function} = nothing`: Optional function that accepts a component
+    of type T and returns a Bool. Apply this function to each component and only return components
+    where the result is true.
+"""
+function get_supplemental_attributes(
+    ::Type{T},
+    component::InfrastructureSystemsComponent,
+) where {T <: SupplementalAttribute}
+    return get_supplemental_attributes(T, get_supplemental_attributes_container(component))
+end
+
+function get_supplemental_attributes(
+    filter_func::Union{Nothing, Function},
+    ::Type{T},
+    component::InfrastructureSystemsComponent,
+) where {T <: SupplementalAttribute}
+    return get_supplemental_attributes(
+        T,
+        get_supplemental_attributes_container(component),
+        filter_func,
+    )
+end
+
+function get_supplemental_attribute(
+    component::InfrastructureSystemsComponent,
+    uuid::Base.UUID,
+)
+    return get_supplemental_attribute(
+        get_supplemental_attributes_container(component),
+        uuid,
+    )
+end
diff --git a/src/components.jl b/src/components.jl
@@ -1,4 +1,3 @@
-
 const ComponentsByType = Dict{DataType, Dict{String, <:InfrastructureSystemsComponent}}
 
 struct Components <: InfrastructureSystemsContainer
@@ -281,31 +280,7 @@ function get_components(
     components::Components,
     filter_func::Union{Nothing, Function} = nothing,
 ) where {T <: InfrastructureSystemsComponent}
-    if isconcretetype(T)
-        _components = get(components.data, T, nothing)
-        if !isnothing(filter_func) && !isnothing(_components)
-            _filter_func = x -> filter_func(x.second)
-            _components = values(filter(_filter_func, _components))
-        end
-        if isnothing(_components)
-            iter = FlattenIteratorWrapper(T, Vector{Base.ValueIterator}([]))
-        else
-            iter =
-                FlattenIteratorWrapper(T, Vector{Base.ValueIterator}([values(_components)]))
-        end
-    else
-        types = [x for x in keys(components.data) if x <: T]
-        if isnothing(filter_func)
-            _components = [values(components.data[x]) for x in types]
-        else
-            _filter_func = x -> filter_func(x.second)
-            _components = [values(filter(_filter_func, components.data[x])) for x in types]
-        end
-        iter = FlattenIteratorWrapper(T, _components)
-    end
-
-    @assert_op eltype(iter) == T
-    return iter
+    return iterate_instances(T, components.data, filter_func)
 end
 
 """

diff --git a/src/geographic_supplemental_attribute.jl b/src/geographic_supplemental_attribute.jl
@@ -1,7 +1,7 @@
 """
 Attribute to store Geographic Information about the system components
 """
-struct GeographicInfo <: InfrastructureSystemsSupplementalAttribute
+struct GeographicInfo <: SupplementalAttribute
     geo_json::Dict{String, Any}
     component_uuids::ComponentUUIDs
     internal::InfrastructureSystemsInternal

diff --git a/src/iterators.jl b/src/iterators.jl
@@ -0,0 +1,31 @@
+function iterate_instances(
+    ::Type{T},
+    data::Union{ComponentsByType, SupplementalAttributesByType},
+    filter_func::Union{Nothing, Function} = nothing,
+) where {T <: InfrastructureSystemsType}
+    if isconcretetype(T)
+        _components = get(data, T, nothing)
+        if !isnothing(filter_func) && !isnothing(_components)
+            _filter_func = x -> filter_func(x.second)
+            _components = values(filter(_filter_func, _components))
+        end
+        if isnothing(_components)
+            iter = FlattenIteratorWrapper(T, Vector{Base.ValueIterator}([]))
+        else
+            iter =
+                FlattenIteratorWrapper(T, Vector{Base.ValueIterator}([values(_components)]))
+        end
+    else
+        types = [x for x in keys(data) if x <: T]
+        if isnothing(filter_func)
+            _components = [values(data[x]) for x in types]
+        else
+            _filter_func = x -> filter_func(x.second)
+            _components = [values(filter(_filter_func, data[x])) for x in types]
+        end
+        iter = FlattenIteratorWrapper(T, _components)
+    end
+
+    @assert_op eltype(iter) == T
+    return iter
+end
diff --git a/src/supplemental_attribute.jl b/src/supplemental_attribute.jl
@@ -1,7 +1,7 @@
 function attach_component!(
     attribute::T,
     component::InfrastructureSystemsComponent,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     component_uuid = get_uuid(component)
 
     if component_uuid ∈ get_component_uuids(attribute)
@@ -17,7 +17,7 @@ function attach_component!(
 end
 
 function detach_component!(
-    attribute::InfrastructureSystemsSupplementalAttribute,
+    attribute::SupplementalAttribute,
     component::InfrastructureSystemsComponent,
 )
     delete!(get_component_uuids(attribute), get_uuid(component))
@@ -27,19 +27,19 @@ end
 """
 Return true if the attribute is attached to at least one component.
 """
-function is_attached_to_component(attribute::InfrastructureSystemsSupplementalAttribute)
+function is_attached_to_component(attribute::SupplementalAttribute)
     return !isempty(get_component_uuids(attribute))
 end
 
 """
 Return true if the attribute has time series data.
 """
-function has_time_series(attribute::InfrastructureSystemsSupplementalAttribute)
+function has_time_series(attribute::SupplementalAttribute)
     container = get_time_series_container(attribute)
     return !isnothing(container) && !isempty(container)
 end
 
-function clear_time_series_storage!(attribute::InfrastructureSystemsSupplementalAttribute)
+function clear_time_series_storage!(attribute::SupplementalAttribute)
     storage = _get_time_series_storage(attribute)
     if !isnothing(storage)
         # In the case of Deterministic and DeterministicSingleTimeSeries the UUIDs
@@ -55,7 +55,7 @@ function clear_time_series_storage!(attribute::InfrastructureSystemsSupplemental
 end
 
 function set_time_series_storage!(
-    attribute::InfrastructureSystemsSupplementalAttribute,
+    attribute::SupplementalAttribute,
     storage::Union{Nothing, TimeSeriesStorage},
 )
     container = get_time_series_container(attribute)
@@ -70,7 +70,7 @@ This function must be called when an attribute is removed from a system.
 """
 function prepare_for_removal!(
     attribute::T,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     if !isempty(get_component_uuids(attribute))
         throw(
             ArgumentError(
@@ -88,7 +88,7 @@ function prepare_for_removal!(
     return
 end
 
-function _get_time_series_storage(attribute::InfrastructureSystemsSupplementalAttribute)
+function _get_time_series_storage(attribute::SupplementalAttribute)
     container = get_time_series_container(attribute)
     if isnothing(container)
         return nothing
@@ -99,7 +99,7 @@ end
 
 function clear_time_series!(
     attribute::T,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     container = get_time_series_container(attribute)
     if !isnothing(container)
         clear_time_series!(container)

diff --git a/src/supplemental_attributes.jl b/src/supplemental_attributes.jl
@@ -1,21 +1,18 @@
-const SupplementalAttributesByType =
-    Dict{DataType, Dict{UUIDs.UUID, <:InfrastructureSystemsSupplementalAttribute}}
-
 struct SupplementalAttributes <: InfrastructureSystemsContainer
-    data::SupplementalAttributesByType
+    data::SupplementalAttributesContainer
     time_series_storage::TimeSeriesStorage
 end
 
 get_member_string(::SupplementalAttributes) = "supplemental attributes"
 
 function SupplementalAttributes(time_series_storage::TimeSeriesStorage)
-    return SupplementalAttributes(SupplementalAttributesByType(), time_series_storage)
+    return SupplementalAttributes(SupplementalAttributesContainer(), time_series_storage)
 end
 
 function add_supplemental_attribute!(
     supplemental_attributes::SupplementalAttributes,
     component::InfrastructureSystemsComponent,
-    supplemental_attribute::InfrastructureSystemsSupplementalAttribute;
+    supplemental_attribute::SupplementalAttribute;
     kwargs...,
 )
     try
@@ -38,7 +35,7 @@ function _add_supplemental_attribute!(
     supplemental_attributes::SupplementalAttributes,
     supplemental_attribute::T;
     allow_existing_time_series = false,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     if !isconcretetype(T)
         throw(ArgumentError("add_supplemental_attribute! only accepts concrete types"))
     end
@@ -116,7 +113,7 @@ end
 function remove_supplemental_attribute!(
     supplemental_attributes::SupplementalAttributes,
     supplemental_attribute::T,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     if !isempty(get_component_uuids(supplemental_attribute))
         throw(
             ArgumentError(
@@ -141,7 +138,7 @@ Throws ArgumentError if the type is not stored.
 function remove_supplemental_attributes!(
     ::Type{T},
     supplemental_attributes::SupplementalAttributes,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
+) where {T <: SupplementalAttribute}
     if !haskey(supplemental_attributes.data, T)
         throw(ArgumentError("supplemental_attribute type $T is not stored"))
     end
@@ -172,50 +169,12 @@ function get_supplemental_attributes(
     ::Type{T},
     supplemental_attributes::SupplementalAttributes,
     filter_func::Union{Nothing, Function} = nothing,
-) where {T <: InfrastructureSystemsSupplementalAttribute}
-    if isconcretetype(T)
-        _supplemental_attributes = get(supplemental_attributes.data, T, nothing)
-        if !isnothing(filter_func) && !isnothing(_supplemental_attributes)
-            _filter_func = x -> filter_func(x.second)
-            _supplemental_attributes =
-                values(filter(_filter_func, _supplemental_attributes))
-        end
-        if isnothing(_supplemental_attributes)
-            iter = FlattenIteratorWrapper(T, Vector{Base.ValueIterator}([]))
-        else
-            iter = FlattenIteratorWrapper(
-                T,
-                Vector{Base.ValueIterator}([values(_supplemental_attributes)]),
-            )
-        end
-    else
-        types = [x for x in keys(supplemental_attributes.data) if x <: T]
-        if isnothing(filter_func)
-            _supplemental_attributes =
-                [values(supplemental_attributes.data[x]) for x in types]
-        else
-            _filter_func = x -> filter_func(x.second)
-            _supplemental_attributes = [
-                values(filter(_filter_func, supplemental_attributes.data[x])) for
-                x in types
-            ]
-        end
-        iter = FlattenIteratorWrapper(T, _supplemental_attributes)
-    end
-
-    @assert_op eltype(iter) == T
-    return iter
+) where {T <: SupplementalAttribute}
+    return get_supplemental_attributes(T, supplemental_attributes.data, filter_func)
 end
 
 function get_supplemental_attribute(attributes::SupplementalAttributes, uuid::Base.UUID)
-    for attr_dict in values(attributes.data)
-        attribute = get(attr_dict, uuid, nothing)
-        if !isnothing(attribute)
-            return attribute
-        end
-    end
-
-    throw(ArgumentError("No attribute with UUID=$uuid is stored"))
+    return get_supplemental_attribute(attributes.data, uuid)
 end
 
 function serialize(attributes::SupplementalAttributes)
@@ -232,7 +191,7 @@ function deserialize(
         type = get_type_from_serialization_metadata(get_serialization_metadata(attr_dict))
         if !haskey(attributes, type)
             attributes[type] =
-                Dict{UUIDs.UUID, InfrastructureSystemsSupplementalAttribute}()
+                Dict{UUIDs.UUID, SupplementalAttribute}()
         end
         attr = deserialize(type, attr_dict)
         uuid = get_uuid(attr)
@@ -243,5 +202,8 @@ function deserialize(
         @debug "Deserialized $(summary(attr))" _group = LOG_GROUP_SERIALIZATION
     end
 
-    return SupplementalAttributes(attributes, time_series_storage)
+    return SupplementalAttributes(
+        SupplementalAttributesContainer(attributes),
+        time_series_storage,
+    )
 end