Merge pull request #527 from LilithHafner/lh/promote

Promote initial state when converting from non-homogeneous python lists.
SciML · Oct 21, 2023 · eab57db · eab57db
2 parents 187e884 + 9cb74b4
commit eab57db
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/ext/SciMLBasePythonCallExt.jl b/ext/SciMLBasePythonCallExt.jl
@@ -13,10 +13,18 @@ function SciMLBase.numargs(f::Py)
     pyconvert(Int, length(first(inspect.getfullargspec(f2))) - inspect.ismethod(f2))
 end
 
-_pyconvert(x::Py) = pyisinstance(x, pybuiltins.list) ? [_pyconvert(x) for x in x] : pyconvert(Any, x)
-_pyconvert(x::PyList) = [_pyconvert(x) for x in x]
+_pyconvert(x::Py) = pyisinstance(x, pybuiltins.list) ? _promoting_collect(_pyconvert(x) for x in x) : pyconvert(Any, x)
+_pyconvert(x::PyList) = _promoting_collect(_pyconvert(x) for x in x)
 _pyconvert(x) = x
 
+# _promoting_collect might copy its input
+_promoting_collect(x) = _promoting_collect(collect(x))
+function _promoting_collect(x::AbstractArray)
+    isconcretetype(eltype(x)) && return x
+    T = mapreduce(typeof, promote_type, x)
+    T == eltype(x) ? x : T.(x)
+end
+
 SciMLBase.prepare_initial_state(u0::Union{Py, PyList}) = _pyconvert(u0)
 SciMLBase.prepare_function(f::Py) = _pyconvert ∘ f
 

diff --git a/test/python/pythoncall.jl b/test/python/pythoncall.jl
@@ -77,3 +77,8 @@ using DifferentialEquations, PythonCall
     sol = de.solve(prob,reltol=1e-3,abstol=1e-3)
     """, @__MODULE__)
 end
+
+@testset "promotion" begin
+    _u0 = pyconvert(Any, pyeval("""de.SciMLBase.prepare_initial_state([1.0, 0, 0])""", @__MODULE__))
+    @test _u0 isa Vector{Float64}
+end