fix custom data (#98)

atoptima · Apr 14, 2023 · bb04c3d · bb04c3d
1 parent fe5d418
commit bb04c3d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/src/callbacks.jl b/src/callbacks.jl
@@ -119,10 +119,6 @@ function MOI.submit(
     return MOI.submit(JuMP.backend(model), cb, JuMP.moi_function(con.func), con.set, custom_data)
 end
 
-MathOptInterface.Utilities.map_indices(
-    variable_map::MathOptInterface.Utilities.IndexMap, x::AbstractCustomData
-) = x
-
 """
 A callback to provide initial columns to the optimizer before starting the optimization.
 """

diff --git a/src/customdata.jl b/src/customdata.jl
@@ -45,10 +45,9 @@ function MOI.set(
     dest::MOIU.UniversalFallback, attribute::CustomVars, value
 )
     if !haskey(dest.modattr, attribute)
-        dest.modattr[attribute] = []
+        dest.modattr[attribute] = Type{<:AbstractCustomData}[]
     end
     for elem in value
-        @assert elem <: AbstractCustomData
         push!(dest.modattr[attribute], elem)
     end
     return
@@ -58,21 +57,20 @@ function MOI.set(
     dest::MOIU.UniversalFallback, attribute::CustomConstrs, value
 )
     if !haskey(dest.modattr, attribute)
-        dest.modattr[attribute] = []
+        dest.modattr[attribute] = Type{<:AbstractCustomData}[]
     end
     for elem in value
-        @assert elem <: AbstractCustomData
         push!(dest.modattr[attribute], elem)
     end
     return
 end
 
 function MOI.get(dest::MOIU.UniversalFallback, attribute::CustomVars)
-    return get(dest.modattr, attribute, [])
+    return get(dest.modattr, attribute, Type{<:AbstractCustomData}[])
 end
 
 function MOI.get(dest::MOIU.UniversalFallback, attribute::CustomConstrs)
-    return get(dest.modattr, attribute, [])
+    return get(dest.modattr, attribute, Type{<:AbstractCustomData}[])
 end
 
 struct CustomVarValue <: MOI.AbstractVariableAttribute end
@@ -149,4 +147,7 @@ end
 
 MathOptInterface.Utilities.map_indices(
     variable_map::MathOptInterface.Utilities.IndexMap, x::AbstractCustomData
+) = x
+MathOptInterface.Utilities.map_indices(
+    variable_map::MathOptInterface.Utilities.IndexMap, x::Vector{AbstractCustomData}
 ) = x