From 57dd17346abbf3365b1e2c8c714efd18a14a9637 Mon Sep 17 00:00:00 2001
From: Dongdong Kong <kongdd@users.noreply.github.com>
Date: Thu, 3 Oct 2024 19:59:09 +0800
Subject: [PATCH] major update for ncvar_def

---
 src/nc_write.jl  | 34 +++++++++----------------------
 src/ncvar_def.jl | 53 ++++++++++++++++++++++++++----------------------
 2 files changed, 39 insertions(+), 48 deletions(-)

diff --git a/src/nc_write.jl b/src/nc_write.jl
index 8e77817..27ad089 100644
--- a/src/nc_write.jl
+++ b/src/nc_write.jl
@@ -44,55 +44,41 @@ $(METHODLIST)
 """
 function nc_write(f::AbstractString, varname::AbstractString, val,
   dims::Vector{NcDim}, attrib::Dict=Dict();
-  units=nothing, longname=nothing,
-  compress=1, overwrite=false, mode="c",
-  global_attrib=Dict(),
-  kw...)
-
-  !isnothing(units) && (attrib["units"] = units)
-  !isnothing(longname) && (attrib["longname"] = longname)
-
+  overwrite=false, mode="c", kw...)
+  
   # check whether variable defined
   if !check_file(f) || overwrite
     isfile(f) && rm(f)
-
-    ds = nc_open(f, mode)
-    ncatt_put(ds, global_attrib)
-    ncdim_def(ds, dims)
-
-    dimnames = names(dims)
-    ncvar_def(ds, varname, val, dimnames, attrib; compress, kw...)
-    close(ds)
+    nc_write!(f, varname, val, dims, attrib; mode, kw...)
   else
     println("[file exist]: $(basename(f))")
   end
 end
 
-
 """
 $(TYPEDSIGNATURES)
 
 $(METHODLIST)
-
-# ! TODO: need to test overwrite
 """
 function nc_write!(f::AbstractString, varname::AbstractString, val,
-  dims::Vector{<:Union{NcDim,AbstractString}}, attrib::Dict=Dict();
+  dims::Vector{<:Union{NcDim,AbstractString}}, attrib::Dict=Dict(); 
   units=nothing, longname=nothing,
-  compress=1, kw...)
+  compress=1, global_attrib=Dict(), mode=nothing, 
+  kw...)
   
   !isnothing(units) && (attrib["units"] = units)
   !isnothing(longname) && (attrib["longname"] = longname)
 
-  mode = check_file(f) ? "a" : "c"
+  isnothing(mode) && (mode = check_file(f) ? "a" : "c")
+
   ds = nc_open(f, mode)
+  ncatt_put(ds, global_attrib)
   ncdim_def(ds, dims; verbose=false)
 
-  ncvar_def(ds, varname, val, dims, attrib; compress=compress, kw...)
+  ncvar_def(ds, varname, val, dims, attrib; compress, kw...)
   close(ds)
 end
 
-
 function nc_write!(f::AbstractString, data::NamedTuple,
   dims::Vector{<:Union{NcDim,AbstractString}}; verbose=false, kw...)
 
diff --git a/src/ncvar_def.jl b/src/ncvar_def.jl
index 44c6d49..a0d852c 100644
--- a/src/ncvar_def.jl
+++ b/src/ncvar_def.jl
@@ -1,8 +1,8 @@
 """
     ncvar_def(ds, name, val, dims::Vector{<:AbstractString}, attrib = Dict();
-        compress = 1, kwargs...)
+        compress = 1, kw...)
     ncvar_def(ds, name, val, dims::NcDim, attrib = Dict();
-        compress = 1, kwargs...)
+        compress = 1, kw...)
 
 # Arguments
 
@@ -57,30 +57,35 @@ nc_info(f)
 
 @seealso [`ncdim_def`](@ref)
 """
-function ncvar_def(ds, name, val, dims::Vector{<:AbstractString}, attrib=Dict();
-    compress=1, type=nothing, overwrite=false, kwargs...)
-
-    # attrib["deflatelevel"] = compress
-    if name in keys(ds) && !overwrite
-        # println(options)
-        @warn "Variable `$name`: exist!"
-        return
+function ncvar_def(ds, name, val, dims::Vector{D}, attrib=Dict();
+  compress=1, type=nothing, overwrite=false, kw...) where {S<:AbstractString,D<:Union{NcDim,S}}
+
+  # change datatype
+  if type !== nothing && eltype(val) != type
+    val = @.(type(val))
+  end
+
+  # attrib["deflatelevel"] = compress
+  if name in keys(ds)
+    if !overwrite
+      @warn "Variable `$name`: exist! Use `overwrite=true` to overwrite."
+      return
+    else
+      @warn "Variable `$name`: exist! Overwrited."
+      
+      _var = ds[name]
+      indices = ntuple(i -> :, ndims(_var))
+      _var[indices...] = val
+      return
     end
-    # change datatype
-    if type !== nothing && eltype(val) != type
-        val = @.(type(val))
-    end
-    # defDim(ds, name, length(val))
-    defVar(ds, name, val, dims; attrib=attrib, deflatelevel=compress, kwargs...)
-end
-
-function ncvar_def(ds, name, val, dims::Vector{NcDim}, attrib=Dict();
-    compress=1, kwargs...)
+  end
 
-    # attrib["deflatelevel"] = compress
+  dimnames = dims
+  # 若只有string name，则无法定义dims，无奈之举
+  if D == NcDim
     ncdim_def(ds, dims) # define dimensions if previous not exist
     dimnames = map(x -> x.name, dims)
-    ncvar_def(ds, name, val, dimnames; attrib=attrib, deflatelevel=compress, kwargs...)
-    # var = NcVar(varname, dims; t = type, compress = compress)
-end
+  end
 
+  defVar(ds, name, val, dimnames; attrib, deflatelevel=compress, kw...)
+end