From c389e3a1039ed66622c6deb243fa5cf3fce8be4b Mon Sep 17 00:00:00 2001
From: Oscar Dowson <odow@users.noreply.github.com>
Date: Wed, 6 Nov 2024 11:19:36 +1300
Subject: [PATCH] [Utilities] ignore attributes in UniversalFallback if set to
 nothing (#2575)

---
 src/Utilities/universalfallback.jl  | 19 +++++++----
 test/Utilities/universalfallback.jl | 52 ++++++++++++++++-------------
 2 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/src/Utilities/universalfallback.jl b/src/Utilities/universalfallback.jl
index 0895206931..4b69179c0f 100644
--- a/src/Utilities/universalfallback.jl
+++ b/src/Utilities/universalfallback.jl
@@ -486,8 +486,10 @@ function MOI.get(
     listattr::MOI.ListOfVariableAttributesSet,
 )
     list = MOI.get(uf.model, listattr)
-    for attr in keys(uf.varattr)
-        push!(list, attr)
+    for (attr, dict) in uf.varattr
+        if !isempty(dict)
+            push!(list, attr)
+        end
     end
     return list
 end
@@ -498,7 +500,7 @@ function MOI.get(
 ) where {F,S}
     list = MOI.get(uf.model, listattr)
     for (attr, dict) in uf.conattr
-        if any(k -> k isa MOI.ConstraintIndex{F,S}, keys(dict))
+        if any(Base.Fix2(isa, MOI.ConstraintIndex{F,S}), keys(dict))
             push!(list, attr)
         end
     end
@@ -694,17 +696,20 @@ function MOI.get(
     )
 end
 
+_set_or_delete(dict, key, value) = (dict[key] = value)
+_set_or_delete(dict, key, ::Nothing) = delete!(dict, key)
+
 function _set(
     uf::UniversalFallback,
     attr::MOI.AbstractOptimizerAttribute,
     value,
 )
-    uf.optattr[attr] = value
+    _set_or_delete(uf.optattr, attr, value)
     return
 end
 
 function _set(uf::UniversalFallback, attr::MOI.AbstractModelAttribute, value)
-    uf.modattr[attr] = value
+    _set_or_delete(uf.modattr, attr, value)
     return
 end
 
@@ -717,7 +722,7 @@ function _set(
     if !haskey(uf.varattr, attr)
         uf.varattr[attr] = Dict{MOI.VariableIndex,Any}()
     end
-    uf.varattr[attr][vi] = value
+    _set_or_delete(uf.varattr[attr], vi, value)
     return
 end
 
@@ -730,7 +735,7 @@ function _set(
     if !haskey(uf.conattr, attr)
         uf.conattr[attr] = Dict{MOI.ConstraintIndex,Any}()
     end
-    uf.conattr[attr][ci] = value
+    _set_or_delete(uf.conattr[attr], ci, value)
     return
 end
 
diff --git a/test/Utilities/universalfallback.jl b/test/Utilities/universalfallback.jl
index 8777c76687..fae6c8df3f 100644
--- a/test/Utilities/universalfallback.jl
+++ b/test/Utilities/universalfallback.jl
@@ -89,21 +89,32 @@ function test_MOI_Test()
     return
 end
 
-function _test_Optimizer_Model_attributes(uf, model, attr, listattr)
+function _test_Optimizer_Model_attributes(
+    uf::MOI.Utilities.UniversalFallback,
+    model::MOI.ModelLike,
+    attr::Union{MOI.AbstractOptimizerAttribute,MOI.AbstractModelAttribute},
+    list::Union{MOI.ListOfOptimizerAttributesSet,MOI.ListOfModelAttributesSet},
+)
     @test !MOI.supports(model, attr)
     @test MOI.supports(uf, attr)
-    @test isempty(MOI.get(uf, listattr))
+    @test isempty(MOI.get(uf, list))
     MOI.set(uf, attr, 0)
     @test MOI.get(uf, attr) == 0
-    @test MOI.get(uf, listattr) == [attr]
+    @test MOI.get(uf, list) == [attr]
+    MOI.set(uf, attr, nothing)
+    @test isempty(MOI.get(uf, list))
+    MOI.set(uf, attr, 0)
     return
 end
 
 function _test_Variable_Constraint_attributes(
-    uf,
-    model,
-    attr,
-    listattr,
+    uf::MOI.Utilities.UniversalFallback,
+    model::MOI.ModelLike,
+    attr::Union{MOI.AbstractVariableAttribute,MOI.AbstractConstraintAttribute},
+    listattr::Union{
+        MOI.ListOfVariableAttributesSet,
+        MOI.ListOfConstraintAttributesSet,
+    },
     I::Type{<:MOI.Index},
     addfun,
     x,
@@ -114,6 +125,11 @@ function _test_Variable_Constraint_attributes(
     @test MOI.supports(uf, attr, I)
     @test isempty(MOI.get(uf, listattr))
     MOI.set(uf, attr, [x, y], [2, 0])
+    @test MOI.get(uf, listattr) == [attr]
+    MOI.set(uf, attr, x, nothing)
+    MOI.set(uf, attr, y, nothing)
+    @test isempty(MOI.get(uf, listattr))
+    MOI.set(uf, attr, [x, y], [2, 0])
     @test MOI.get(uf, attr, z) === nothing
     @test !MOI.is_empty(uf)
     @test MOI.get(uf, listattr) == [attr]
@@ -235,18 +251,13 @@ function test_supported_constraint_attributes()
     cx = _add_constraint(uf, x, 0.0)
     cy = _add_constraint(uf, y, 1.0)
     cz = _add_constraint(uf, z, 2.0)
+    F, S = MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}
     _test_Variable_Constraint_attributes(
         uf,
         model,
         MOI.Test.UnknownConstraintAttribute(),
-        MOI.ListOfConstraintAttributesSet{
-            MOI.ScalarAffineFunction{Float64},
-            MOI.LessThan{Float64},
-        }(),
-        MOI.ConstraintIndex{
-            MOI.ScalarAffineFunction{Float64},
-            MOI.LessThan{Float64},
-        },
+        MOI.ListOfConstraintAttributesSet{F,S}(),
+        MOI.ConstraintIndex{F,S},
         uf -> _add_constraint(uf, x, 0.0),
         cx,
         cy,
@@ -282,18 +293,13 @@ function test_unsupported_constraint_attributes()
     cx = _add_constraint(uf, x, 0.0)
     cy = _add_constraint(uf, y, 1.0)
     cz = _add_constraint(uf, z, 2.0)
+    F, S = MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}
     _test_Variable_Constraint_attributes(
         uf,
         model,
         MOI.Test.UnknownConstraintAttribute(),
-        MOI.ListOfConstraintAttributesSet{
-            MOI.ScalarAffineFunction{Float64},
-            MOI.EqualTo{Float64},
-        }(),
-        MOI.ConstraintIndex{
-            MOI.ScalarAffineFunction{Float64},
-            MOI.EqualTo{Float64},
-        },
+        MOI.ListOfConstraintAttributesSet{F,S}(),
+        MOI.ConstraintIndex{F,S},
         uf -> _add_constraint(uf, x, 0.0),
         cx,
         cy,