diff --git a/ext/AwkwardPythonCallExt/AwkwardPythonCallExt.jl b/ext/AwkwardPythonCallExt/AwkwardPythonCallExt.jl index fed937f..eab0a4f 100644 --- a/ext/AwkwardPythonCallExt/AwkwardPythonCallExt.jl +++ b/ext/AwkwardPythonCallExt/AwkwardPythonCallExt.jl @@ -32,4 +32,52 @@ function AwkwardArray.convert(array::Py)::AwkwardArray.Content ) end +function AwkwardArray.convert(array::PyIterable{Any})::AwkwardArray.Content + form, len, _containers = pyimport("awkward").to_buffers(array) + containers = pyconvert(Dict, _containers) + + julia_buffers = Dict{String,AbstractVector{UInt8}}() + for (key, buffer) in containers + julia_buffers[key] = reinterpret(UInt8, buffer) + end + + AwkwardArray.from_buffers( + pyconvert(String, form.to_json()), + pyconvert(Int, len), + julia_buffers, + ) +end + + +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.PrimitiveArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.EmptyArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.ListOffsetArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.ListArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.RegularArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.RecordArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.TupleArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.IndexedArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.IndexedOptionArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.ByteMaskedArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.BitMaskedArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.UnmaskedArray}, x::Py) = AwkwardArray.convert(x) +PythonCall.pyconvert_rule_array_nocopy(::Type{AwkwardArray.UnionArray}, x::Py) = AwkwardArray.convert(x) + + +function __init__() + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.PrimitiveArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.EmptyArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ListOffsetArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ListArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.RegularArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.RecordArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.TupleArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.IndexedArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.IndexedOptionArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.ByteMaskedArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.BitMaskedArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.UnmaskedArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) + PythonCall.pyconvert_add_rule("awkward.highlevel:Array", AwkwardArray.UnionArray, PythonCall.pyconvert_rule_array_nocopy, PythonCall.PYCONVERT_PRIORITY_ARRAY) +end + end # module diff --git a/test/runpytests.jl b/test/runpytests.jl index e0ab499..fd65e29 100644 --- a/test/runpytests.jl +++ b/test/runpytests.jl @@ -30,3 +30,15 @@ end @test array == [[1.1, 2.2, 3.3], [], [4.4, 5.5]] end + +# Test passing Python array to Julia function +@testset "pass Python array to Julia test" begin + function f1(x) + x + end + + py_array = pyimport("awkward").Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]) + + array = f1(py_array) + @test array isa AwkwardArray.ListOffsetArray +end