fix: cleanup and rename convert functions

JuliaHEP · Nov 28, 2023 · 9453842 · 9453842
1 parent 606b531
commit 9453842
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 29 deletions.
diff --git a/src/pycall/AwkwardPyCall.jl b/src/pycall/AwkwardPyCall.jl
@@ -13,37 +13,29 @@ end
 using PyCall
 using AwkwardArray
 
-function _as_numpy(array::AbstractVector{UInt8})
-    py_array = PyObject(array)
-    pyimport("numpy").asarray(py_array, dtype = pyimport("numpy").uint8)
-end
-
-function julia_array_to_python(array)
-    form, len, containers = AwkwardArray.to_buffers(array)
+function convert(layout::AwkwardArray.Content)::PyObject
+    form, len, containers = AwkwardArray.to_buffers(layout)
 
     py_buffers = Dict{String,Any}()
 
     for (key, buffer) in containers
-        py_buffers[key] = _as_numpy(buffer)
+        py_buffers[key] =
+            pyimport("numpy").asarray(PyObject(buffer), dtype = pyimport("numpy").uint8)
     end
 
-    return pyimport("awkward").from_buffers(form, len, py_buffers)
+    pyimport("awkward").from_buffers(form, len, py_buffers)
 end
 
-function _as_julia(py_buffer)
-    uint8_buffer = reinterpret(UInt8, py_buffer)
-    return uint8_buffer
-end
-
-function python_array_to_julia(py_array)
-    form, len, containers = pyimport("awkward").to_buffers(py_array)
+function convert(array::PyObject)::AwkwardArray.Content
+    form, len, containers = pyimport("awkward").to_buffers(array)
 
     julia_buffers = Dict{String,AbstractVector{UInt8}}()
+
     for (key, buffer) in containers
-        julia_buffers[key] = _as_julia(buffer)
+        julia_buffers[key] = reinterpret(UInt8, buffer)
     end
 
-    return AwkwardArray.from_buffers(form.to_json(), len, julia_buffers)
+    AwkwardArray.from_buffers(form.to_json(), len, julia_buffers)
 end
 
 end # module
diff --git a/test/runpytests.jl b/test/runpytests.jl
@@ -22,34 +22,32 @@ using JSON
 
 include("../src/pycall/AwkwardPyCall.jl")
 
-import Main.AwkwardPyCall: julia_array_to_python
+import Main.AwkwardPyCall: convert
 
-# Test julia_array_to_python function
-@testset "julia_array_to_python tests" begin
+# Test convert Julia array to Python function
+@testset "convert Julia array to Python tests" begin
     array = AwkwardArray.ListOffsetArray(
         [0, 3, 3, 5],
         AwkwardArray.PrimitiveArray([1.1, 2.2, 3.3, 4.4, 5.5]),
     )
     # Test case 1: Check if the function returns an awkward array
-    @test isa(julia_array_to_python(array), PyObject)
+    @test isa(convert(array), PyObject)
 
     # Test case 2: Check if the awkward array has the correct layout
-    py_array = julia_array_to_python(array)
+    py_array = convert(array)
     @test typeof(py_array) == PyObject
     @test pyimport("awkward").to_list(py_array) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
 
 
 end
 
-import Main.AwkwardPyCall: python_array_to_julia
-
-# Test python_array_to_julia function
-@testset "python_array_to_julia tests" begin
+# Test convert Python array to Julia function
+@testset "convert Python array to Julia tests" begin
     py_array = pyimport("awkward").Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
 
     # Test case 1: Check if the function returns an awkward array
     @test isa(
-        python_array_to_julia(py_array),
+        convert(py_array),
         AwkwardArray.ListOffsetArray{
             SubArray{Int64,1,Vector{Int64},Tuple{UnitRange{Int64}},true},
             AwkwardArray.PrimitiveArray{
@@ -62,7 +60,7 @@ import Main.AwkwardPyCall: python_array_to_julia
     )
 
     # Test case 2: Check if the awkward array has the correct layout
-    array = python_array_to_julia(py_array)
+    array = convert(py_array)
     @test typeof(array) == AwkwardArray.ListOffsetArray{
         SubArray{Int64,1,Vector{Int64},Tuple{UnitRange{Int64}},true},
         AwkwardArray.PrimitiveArray{