add jump methods

src/jump_methods.jl

@@ -1,3 +1,12 @@
+"""
+    JuMP.constraint_ref_with_index(
+    element::OptiElement, 
+    idx::MOI.ConstraintIndex{<:MOI.AbstractScalarFunction, <:MOI.AbstractScalarSet}
+    )
+
+Return a `ConstraintRef` given an optigraph element and `MOI.ConstraintIndex`. 
+Note that the index is the index corresponding to the graph backend, not the element index.
+"""
 function JuMP.constraint_ref_with_index(
     element::OptiElement, 
     idx::MOI.ConstraintIndex{<:MOI.AbstractScalarFunction, <:MOI.AbstractScalarSet}
@@ -9,11 +18,6 @@ function JuMP.constraint_ref_with_index(element::OptiElement, idx::MOI.VariableI
     return JuMP.constraint_ref_with_index(graph_backend(element), idx)
 end
 
-function JuMP.list_of_constraint_types(graph::OptiGraph)::Vector{Tuple{Type,Type}}
-    all_constraint_types = JuMP.list_of_constraint_types.(all_nodes(graph))
-    return unique(vcat(all_constraint_types...))
-end
-
 function JuMP.list_of_constraint_types(element::OptiElement)::Vector{Tuple{Type,Type}}
     # NOTE from JuMP:
     # We include an annotated return type here because Julia fails terribly at
@@ -24,20 +28,40 @@ function JuMP.list_of_constraint_types(element::OptiElement)::Vector{Tuple{Type,
     ]
 end
 
+function JuMP.list_of_constraint_types(graph::OptiGraph)::Vector{Tuple{Type,Type}}
+    all_constraint_types = JuMP.list_of_constraint_types.(all_elements(graph))
+    return unique(vcat(all_constraint_types...))
+end
+
+### num_constraints
+
+"""
+    JuMP.num_constraints(
+    element::OptiElement,
+    function_type::Type{
+        <:Union{JuMP.AbstractJuMPScalar,Vector{<:JuMP.AbstractJuMPScalar}},
+    },set_type::Type{<:MOI.AbstractSet})::Int64
+
+Return the total number of constraints on an element.
+"""
 function JuMP.num_constraints(
     element::OptiElement,
     function_type::Type{
         <:Union{JuMP.AbstractJuMPScalar,Vector{<:JuMP.AbstractJuMPScalar}},
     },
-    set_type::Type{<:MOI.AbstractSet},
+    set_type::Type{<:MOI.AbstractSet}
 )::Int64
     F = JuMP.moi_function_type(function_type)
     return MOI.get(element, MOI.NumberOfConstraints{F,set_type}())
 end
 
 function JuMP.num_constraints(graph::OptiGraph)
+    all_num_constraints = JuMP.num_constraints.(all_elements(graph))
+    return sum(all_num_constraints)
 end
 
+### all_constraints
+
 function JuMP.all_constraints(
     element::OptiElement,
     func_type::Type{
@@ -63,5 +87,18 @@ function JuMP.all_constraints(
     return result
 end
 
-function JuMP.all_constraints(graph::OptiGraph)
+# return all of the constraints, including constraints in subgraphs
+function JuMP.all_constraints(
+    graph::OptiGraph,
+    func_type::Type{
+        <:Union{JuMP.AbstractJuMPScalar,Vector{<:JuMP.AbstractJuMPScalar}},
+    },
+    set_type::Type{<:MOI.AbstractSet})
+
+    all_graph_constraints = JuMP.all_constraints.(
+        all_elements(graph), 
+        Ref(func_type),
+        Ref(set_type)
+    )
+    return vcat(all_graph_constraints...)
 end

src/macros.jl

@@ -72,7 +72,6 @@ end
     @nodevariables(iterable, expr...)
 
 Call the JuMP.@variable macro for each optinode in a given container
-
 """
 macro nodevariables(nodes, args...)
     macro_code = quote

src/optigraph.jl

@@ -195,8 +195,6 @@ function JuMP.all_variables(graph::OptiGraph)
     return vcat(JuMP.all_variables.(all_nodes(graph))...)
 end
 
-
-
 ### OptiEdges
 
 function get_edge(graph::OptiGraph, nodes::Set{OptiNode})
@@ -251,6 +249,9 @@ function all_edges(graph::OptiGraph)
     return edges
 end
 
+function all_elements(graph::OptiGraph)
+    return [all_nodes(graph); all_edges(graph)]
+end
 
 ### MOI Methods
 

src/optinode.jl

@@ -56,8 +56,6 @@ function graph_backend(node::OptiNode)
     return graph_backend(source_graph(node))
 end
 
-
-
 """
     Filter the object dictionary for values that belong to node. Keep in mind that 
 this function is slow for optigraphs with many nodes.
@@ -241,7 +239,7 @@ end
 
 function MOI.set(
     node::OptiNode,
-    attr::MOI.AbstractConstraintAttribute,
+    attr::AT,
     cref::ConstraintRef,
     args...
 ) where AT <: MOI.AbstractConstraintAttribute