refactor: distinguish unloaded extensions from dispatch errors

ext/DynamicExpressionsBumperExt.jl

@@ -5,9 +5,10 @@ using DynamicExpressions:
     OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalContext
 using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
-import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array, bumper_kern!
+import DynamicExpressions.ExtensionInterfaceModule:
+    _bumper_eval_tree_array, is_extension_loaded, bumper_kern!
 
-function bumper_eval_tree_array(
+function _bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
@@ -87,4 +88,6 @@ function bumper_kern!(
     return cumulator_1
 end
 
+is_extension_loaded(::Val{:Bumper}) = true
+
 end

ext/DynamicExpressionsSymbolicUtilsExt.jl

@@ -10,7 +10,12 @@ using DynamicExpressions.UtilsModule: deprecate_varmap
 using SymbolicUtils
 using SymbolicUtils: BasicSymbolic, SymReal, iscall, issym, isconst, unwrap_const
 
-import DynamicExpressions.ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
+import DynamicExpressions.ExtensionInterfaceModule:
+    is_extension_loaded,
+    _node_to_symbolic,
+    _symbolic_to_node,
+    node_to_symbolic,
+    symbolic_to_node
 import DynamicExpressions.ValueInterfaceModule: is_valid
 
 const SYMBOLIC_UTILS_TYPES = Union{<:Number,BasicSymbolic}
@@ -268,25 +273,7 @@ function Base.convert(
     return constructorof(E)(tree; operators, variable_names, kws...)
 end
 
-"""
-    node_to_symbolic(tree::AbstractExpressionNode, operators::AbstractOperatorEnum;
-                variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing,
-                index_functions::Bool=false)
-
-The interface to SymbolicUtils.jl. Passing a tree to this function
-will generate a symbolic equation in SymbolicUtils.jl format.
-
-## Arguments
-
-- `tree::AbstractExpressionNode`: The equation to convert.
-- `operators::AbstractOperatorEnum`: OperatorEnum, which contains the operators used in the equation.
-- `variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing`: What variable names to use for
-    each feature. Default is [x1, x2, x3, ...].
-- `index_functions::Bool=false`: Whether to generate special names for the
-    operators, which then allows one to convert back to a `AbstractExpressionNode` format
-    using `symbolic_to_node`.
-"""
-function node_to_symbolic(
+function _node_to_symbolic(
     tree::AbstractExpressionNode{T,2},
     operators::AbstractOperatorEnum;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
@@ -309,7 +296,7 @@ function node_to_symbolic(
     )
     return substitute(expr, subs)
 end
-function node_to_symbolic(
+function _node_to_symbolic(
     tree::AbstractExpression,
     operators::Union{AbstractOperatorEnum,Nothing}=nothing;
     variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing,
@@ -323,7 +310,7 @@ function node_to_symbolic(
     )
 end
 
-function symbolic_to_node(
+function _symbolic_to_node(
     eqn::BasicSymbolic,
     operators::AbstractOperatorEnum,
     ::Type{N}=Node;
@@ -429,4 +416,6 @@ function multiply_powers(
     end
 end
 
+is_extension_loaded(::Val{:SymbolicUtils}) = true
+
 end

ext/DynamicExpressionsZygoteExt.jl

@@ -1,9 +1,10 @@
 module DynamicExpressionsZygoteExt
 
 using Zygote: gradient
-import DynamicExpressions.ExtensionInterfaceModule: _zygote_gradient, ZygoteGradient
+import DynamicExpressions.ExtensionInterfaceModule:
+    is_extension_loaded, _zygote_gradient_impl, ZygoteGradient
 
-function _zygote_gradient(op::F, ::Val{degree}) where {F,degree}
+function _zygote_gradient_impl(op::F, ::Val{degree}) where {F,degree}
     return ZygoteGradient{F,degree}(op)
 end
 
@@ -13,4 +14,6 @@ function (g::ZygoteGradient{F,degree})(args::Vararg{Any,degree}) where {F,degree
     return ntuple(i -> @something(partials[i], zero(args[i])), Val(degree))
 end
 
+is_extension_loaded(::Val{:Zygote}) = true
+
 end

src/ExtensionInterface.jl

@@ -1,11 +1,40 @@
 module ExtensionInterfaceModule
 
-function node_to_symbolic(args...; kws...)
-    return error("Please load the `SymbolicUtils` package to use `node_to_symbolic`.")
+using DispatchDoctor: @unstable
+
+is_extension_loaded(::Val) = false
+
+"""
+    node_to_symbolic(tree::AbstractExpressionNode, operators::AbstractOperatorEnum;
+                variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing,
+                index_functions::Bool=false)
+
+The interface to SymbolicUtils.jl. Passing a tree to this function
+will generate a symbolic equation in SymbolicUtils.jl format.
+
+## Arguments
+
+- `tree::AbstractExpressionNode`: The equation to convert.
+- `operators::AbstractOperatorEnum`: OperatorEnum, which contains the operators used in the equation.
+- `variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing`: What variable names to use for
+    each feature. Default is [x1, x2, x3, ...].
+- `index_functions::Bool=false`: Whether to generate special names for the
+    operators, which then allows one to convert back to a `AbstractExpressionNode` format
+    using `symbolic_to_node`.
+"""
+@unstable function node_to_symbolic(args...; kws...)
+    is_extension_loaded(Val(:SymbolicUtils)) ||
+        error("Please load the `SymbolicUtils` package to use `node_to_symbolic`.")
+    return _node_to_symbolic(args...; kws...)
 end
-function symbolic_to_node(args...; kws...)
-    return error("Please load the `SymbolicUtils` package to use `symbolic_to_node`.")
+function _node_to_symbolic end
+
+@unstable function symbolic_to_node(args...; kws...)
+    is_extension_loaded(Val(:SymbolicUtils)) ||
+        error("Please load the `SymbolicUtils` package to use `symbolic_to_node`.")
+    return _symbolic_to_node(args...; kws...)
 end
+function _symbolic_to_node end
 
 struct ZygoteGradient{F,degree} <: Function
     op::F
@@ -19,12 +48,17 @@ end
 Base.show(io::IO, ::MIME"text/plain", g::ZygoteGradient) = show(io, g)
 
 function _zygote_gradient(args...)
-    return error("Please load the Zygote.jl package.")
+    is_extension_loaded(Val(:Zygote)) || error("Please load the Zygote.jl package.")
+    return _zygote_gradient_impl(args...)
 end
+function _zygote_gradient_impl end
 
 function bumper_eval_tree_array(args...)
-    return error("Please load the Bumper.jl package to use this feature.")
+    is_extension_loaded(Val(:Bumper)) ||
+        error("Please load the Bumper.jl package to use this feature.")
+    return _bumper_eval_tree_array(args...)
 end
+function _bumper_eval_tree_array end
 function bumper_kern! end
 
 _is_loopvectorization_loaded(_) = false  # COV_EXCL_LINE

test/test_initial_errors.jl

@@ -1,5 +1,7 @@
 using DynamicExpressions
 using DynamicExpressions: EvalContext
+using DynamicExpressions.ExtensionInterfaceModule:
+    is_extension_loaded, _zygote_gradient, bumper_eval_tree_array
 using DispatchDoctor: allow_unstable
 using Test
 
@@ -46,3 +48,41 @@ tree = cos(2.1 * x1) + sin(x2)
         () -> tree(ones(2, 10), operators; eval_context=EvalContext(; turbo=Val(true)))
     )
 )
+
+# Loaded extensions should use normal dispatch for unsupported arguments instead of claiming
+# that the dependency is missing.
+using SymbolicUtils
+
+@test is_extension_loaded(Val(:SymbolicUtils))
+@test !is_extension_loaded(Val(:Zygote))
+@test !is_extension_loaded(Val(:Bumper))
+@test_throws "Please load the Zygote.jl package." _zygote_gradient(nothing)
+@test_throws "Please load the Bumper.jl package" bumper_eval_tree_array(nothing)
+
+symbolic_x1 = allow_unstable(() -> node_to_symbolic(x1, operators))
+@test string(symbolic_x1) == "x1"
+@test string(allow_unstable(() -> symbolic_to_node(symbolic_x1, operators))) == "x1"
+@test_throws MethodError node_to_symbolic(nothing)
+@test_throws MethodError symbolic_to_node(nothing)
+
+using Zygote
+
+@test is_extension_loaded(Val(:SymbolicUtils))
+@test is_extension_loaded(Val(:Zygote))
+@test !is_extension_loaded(Val(:Bumper))
+@test_throws "Please load the Bumper.jl package" bumper_eval_tree_array(nothing)
+
+@test allow_unstable(() -> only(_zygote_gradient(sin, Val(1))(1.0))) ≈ cos(1.0)
+@test_throws MethodError _zygote_gradient(nothing)
+
+using Bumper
+
+@test is_extension_loaded(Val(:SymbolicUtils))
+@test is_extension_loaded(Val(:Zygote))
+@test is_extension_loaded(Val(:Bumper))
+
+bumper_result = allow_unstable(
+    () -> tree(ones(2, 10), operators; eval_options=EvalContext(; bumper=Val(true)))
+)
+@test bumper_result ≈ fill(cos(2.1) + sin(1.0), 10)
+@test_throws MethodError bumper_eval_tree_array(nothing)