From 350f6ed19ead4c31fee69df779b1aea841e368ad Mon Sep 17 00:00:00 2001
From: Dongdong Kong <kongdd@users.noreply.github.com>
Date: Sun, 17 Dec 2023 22:13:20 +0800
Subject: [PATCH] improve ncdim_def

---
 src/nc_write.jl  |  2 ++
 src/ncdim_def.jl | 18 ++++++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/nc_write.jl b/src/nc_write.jl
index a5f438a..251eb37 100644
--- a/src/nc_write.jl
+++ b/src/nc_write.jl
@@ -87,6 +87,8 @@ function nc_write!(f::AbstractString, varname::AbstractString, val, dims::Vector
 
   mode = check_file(f) ? "a" : "c"
   ds = nc_open(f, mode)
+  ncdim_def(ds, dims; verbose=false)
+  
   ncvar_def(ds, varname, val, dims, attrib; compress=compress, kw...)
   close(ds)
 end
diff --git a/src/ncdim_def.jl b/src/ncdim_def.jl
index 041aa34..2195385 100644
--- a/src/ncdim_def.jl
+++ b/src/ncdim_def.jl
@@ -1,5 +1,9 @@
-# TODO: 这里有提升空间
-function NcDim_time(dates)
+import Dates: Date, DateTime
+import CFTime: AbstractCFDateTime
+
+const DateTimeType = Union{DateTime,AbstractCFDateTime}
+
+function NcDim_time(dates::Vector{<:DateTimeType})
   vals = CFTime.timeencode.(dates, "days since 1970-01-01", eltype(dates))
   attrib = Dict(
     "units" => "days since 1970-01-01",
@@ -8,6 +12,8 @@ function NcDim_time(dates)
   NcDim("time", length(vals), vals, attrib)
 end
 
+NcDim_time(dates::Vector{Date}) = NcDim_time(DateTime.(dates))
+
 """
   make_dims(range=[70, 140, 15, 55], cellsize = 0.5, dates)
   
@@ -60,13 +66,13 @@ function ncdim_def(ds, name, val, attrib=Dict(); verbose=false)
 end
 
 # NcDim
-function ncdim_def(ds, dim::NcDim)
-  ncdim_def(ds, dim.name, dim.vals, dim.atts)
+function ncdim_def(ds, dim::NcDim; kw...)
+  ncdim_def(ds, dim.name, dim.vals, dim.atts; kw...)
 end
 
-function ncdim_def(ds, dims::Vector{NcDim})
+function ncdim_def(ds, dims::Vector{NcDim}; kw...)
   for i = 1:length(dims)
     dim = dims[i]
-    ncdim_def(ds, dim.name, dim.vals, dim.atts)
+    ncdim_def(ds, dim.name, dim.vals, dim.atts; kw...)
   end
 end