From 0c58eb84abf5619c3f1f3dc0dce49541352e4332 Mon Sep 17 00:00:00 2001
From: Beforerr <zzj956959688@gmail.com>
Date: Mon, 3 Feb 2025 13:59:46 -0800
Subject: [PATCH] feat: convert VariableAxis

---
 ext/DimensionalDataExt.jl |  6 +++++-
 src/Speasy.jl             | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/ext/DimensionalDataExt.jl b/ext/DimensionalDataExt.jl
index 667abb0..f151c79 100644
--- a/ext/DimensionalDataExt.jl
+++ b/ext/DimensionalDataExt.jl
@@ -1,14 +1,18 @@
 module DimensionalDataExt
 using DimensionalData
 using Speasy
+using Speasy: axes
 using Unitful
 import DimensionalData: DimArray, DimStack
 
 function DimArray(v::SpeasyVariable; unit=unit(v))
     v = replace_fillval_by_nan(v)
+    axes = v.axes
     name = Symbol(v.name)
     dims = (Ti(v.time), Dim{name}(v.columns))
-    DimArray(v.values * unit, dims; name, metadata=v.meta)
+    metadata = v.meta
+    metadata["axes"] = axes
+    DimArray(v.values * unit, dims; name, metadata)
 end
 
 function DimArray(vs::AbstractArray{SpeasyVariable})
diff --git a/src/Speasy.jl b/src/Speasy.jl
index 79291cb..5d72f20 100644
--- a/src/Speasy.jl
+++ b/src/Speasy.jl
@@ -24,6 +24,14 @@ struct SpeasyVariable
     py::Py
 end
 
+"""
+A wrapper of `speasy.VariableAxis`.
+https://github.com/SciQLop/speasy/blob/main/speasy/core/data_containers.py#L229
+"""
+struct VariableAxis
+    py::Py
+end
+
 function get_data(args...)
     res = @pyconst(pyimport("speasy").get_data)(args...)
     return apply_recursively(res, SpeasyVariable, is_pylist)
@@ -35,11 +43,24 @@ getindex(var::SpeasyVariable, s::Symbol) = getindex(var, string(s))
 name(var) = pyconvert(String, var.py.name)
 values(var) = pyconvert(Array, var.py.values)
 time(var) = pyconvert_time(var.py.time)
+axes(var, i) = VariableAxis(var.py.axes[i-1])
+axes(var) = [axes(var, i) for i in 1:pylen(var.py.axes)]
 columns(var) = pyconvert(Vector{Symbol}, var.py.columns)
 meta(var) = pyconvert(Dict, var.py.meta)
 units(var) = pyconvert(String, var.py.unit)
 coord(var) = pyconvert(String, var.py.meta["COORDINATE_SYSTEM"])
 
+function values(ax::VariableAxis)
+    if ax.name == "time"
+        pyconvert_time(ax.py.values)
+    else
+        pyconvert(Array, ax.py.values)
+    end
+end
+
+units(ax::VariableAxis) = pyconvert(String, ax.py.unit)
+name(ax::VariableAxis) = pyconvert(String, ax.py.name)
+
 function Unitful.unit(var::SpeasyVariable)
     u_str = units(var)
     try
@@ -55,7 +76,7 @@ function Unitful.unit(var::SpeasyVariable)
     return 1
 end
 
-const speasy_properties = (:name, :values, :time, :columns, :meta, :units)
+const speasy_properties = (:name, :values, :time, :columns, :meta, :units, :axes)
 
 function getproperty(var::SpeasyVariable, s::Symbol)
     s in (:py,) && return getfield(var, s)
@@ -65,6 +86,16 @@ end
 
 propertynames(var::SpeasyVariable) = union(fieldnames(SpeasyVariable), speasy_properties)
 
+ax_properties = (:name, :values, :units)
+
+function getproperty(var::VariableAxis, s::Symbol)
+    s in (:py,) && return getfield(var, s)
+    s in ax_properties && return eval(s)(var)
+    return getproperty(var.py, s)
+end
+
+propertynames(var::VariableAxis) = union(fieldnames(VariableAxis), ax_properties)
+
 function speasyplot end
 function speasyplot! end