Merge pull request #35 from sintefmath/jutulstorage_type

Fix type stability issues for JutulStorage
sintefmath · Jun 27, 2023 · 4884920 · 4884920
2 parents 9666bc4 + a71367a
commit 4884920
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/core_types/core_types.jl b/src/core_types/core_types.jl
@@ -487,7 +487,7 @@ end
 import Base: getindex, @propagate_inbounds, parent, size, axes
 
 struct JutulStorage{K}
-    data::Union{Dict{Symbol, <:Any}, NamedTuple}
+    data::Union{Dict{Symbol, Any}, K}
     function JutulStorage(S = Dict{Symbol, Any}(); kwarg...)
         if isa(S, Dict)
             K = Nothing
@@ -496,8 +496,7 @@ struct JutulStorage{K}
             end
         else
             @assert isa(S, NamedTuple)
-            K = keys(S)
-            K::Tuple
+            K = typeof(S)
             @assert length(kwarg) == 0
         end
         return new{K}(S)
@@ -523,7 +522,7 @@ end
 Base.propertynames(S::JutulStorage) = keys(getfield(S, :data))
 
 data(S::JutulStorage{Nothing}) = getfield(S, :data)
-data(S::JutulStorage) = getfield(S, :data)::NamedTuple
+data(S::JutulStorage{T}) where T = getfield(S, :data)::T
 
 function Base.setproperty!(S::JutulStorage, name::Symbol, x)
     Base.setproperty!(data(S), name, x)