carray-netcdf 1.0.13 → 2.0.0

data/lib/carray/netcdf/file.rb

@@ -0,0 +1,354 @@
+module CA
+  module NetCDF
+    # A NetCDF file — unified read / write / create.
+    #
+    # @example
+    #   File.open("data.nc")                  # read
+    #   File.open("data.nc", :write)          # modify existing
+    #   File.create("out.nc")                 # new file
+    #   File.create("out.nc", atomic: true)   # tmp + rename on success
+    #
+    # The `File` IS the root {Group}; all Group methods (`dims`, `vars`,
+    # `def_dim`, `def_var`, `def_group`, `nc[name]`, `nc[name]=`, …) work
+    # directly on it and operate on the root group. Named subgroups are
+    # {Group} instances reachable via {#groups} / {Group#group} /
+    # {Group#def_group}.
+    #
+    # Block forms auto-close (and atomic writes commit at that point).
+    # Names are Strings throughout.
+    class File < Group
+
+      # Kwargs on {.create} that become file-wide defaults for every
+      # subsequent {Group#def_var} call. See docs/writing.md for details.
+      # @return [Array<Symbol>]
+      FILE_VAR_DEFAULTS = %i[chunks deflate shuffle fletcher32 endian no_fill].freeze
+
+      class << self
+        # Open an existing NetCDF file.
+        #
+        # @param path [String]
+        # @param mode [:read, :write, :append]
+        # @param codec [Symbol, Array, Object] see docs/codec.md
+        # @yieldparam file [File]
+        # @return [File] the file (or block's return value if a block is
+        #   given; the file is closed after the block)
+        def open(path, mode = :read, codec: :auto, &block)
+          flag = case mode
+                 when :read           then ::NetCDF::NC_NOWRITE
+                 when :write, :append then ::NetCDF::NC_WRITE | ::NetCDF::NC_SHARE
+                 else raise ArgumentError, "unknown mode: #{mode.inspect}"
+                 end
+          build(::NetCDF.nc_open(path, flag), path,
+                mode: mode, codec: codec, defining: false, &block)
+        end
+
+        # Create a new NetCDF file.
+        #
+        # @param path [String] final file path
+        # @param codec [Symbol, Array, Object] see docs/codec.md
+        # @param clobber [true, false, :skip] overwrite policy when `path`
+        #   already exists. `true` (default) overwrites; `false` raises
+        #   `Errno::EEXIST`; `:skip` returns `nil` without calling the
+        #   block — handy for resumable batches (`Dir[...]` picks up
+        #   completed outputs unchanged).
+        # @param netcdf4 [Boolean] use the NetCDF-4 (HDF5) format. When a
+        #   `template:` File / Group is provided, this defaults to
+        #   whatever the template's format was, so a NetCDF-4 template
+        #   produces a NetCDF-4 output unless explicitly overridden.
+        # @param atomic [Boolean] write to a same-dir tmp file, rename to
+        #   `path` only on successful close; unlink the tmp on any
+        #   exception from the block form. Rename is atomic on POSIX.
+        #   Consumers polling the target path won't see anything until
+        #   close — for a long-running unlimited-dim streaming write,
+        #   prefer `atomic: false`.
+        # @param template [File, Group, nil] a reference file / group
+        #   whose skeleton (dimensions and global attributes) is copied
+        #   into the new file before the block runs. See `dims:`, `atts:`,
+        #   `vars:`.
+        # @param dims [true, false, Array<String>] which template dims to
+        #   copy. `true` (default) = all dims; `false` = none; Array =
+        #   named subset. Ignored when `template:` is nil.
+        # @param atts [true, false, Hash] which template global attributes
+        #   to copy. `true` (default) = all; `false` = none; Hash =
+        #   template's atts merged with these overrides (last-writer
+        #   wins).
+        # @param vars [Array<String>] template variable names to copy in
+        #   full (definition AND data, bit-identical via raw + layer 1).
+        #   Typical use: copy coordinate vars like `%w[time lat lon]` once
+        #   per batch iteration.
+        # @param var_defaults [Hash] file-wide defaults picked up by every
+        #   subsequent `def_var` (a per-var kwarg on `def_var` overrides
+        #   the file default). Valid keys: {FILE_VAR_DEFAULTS}.
+        # @yieldparam file [File]
+        # @return [File, nil] the file (or block's return value); `nil`
+        #   when `clobber: :skip` and `path` already exists.
+        def create(path, codec: :auto, clobber: true, netcdf4: false,
+                   atomic: false, template: nil,
+                   dims: true, atts: true, vars: [],
+                   **var_defaults, &block)
+          # Extract file-wide var defaults; anything else is an error.
+          unknown = var_defaults.keys - FILE_VAR_DEFAULTS
+          unless unknown.empty?
+            raise ArgumentError, "unknown File.create kwargs: #{unknown.inspect}"
+          end
+          # A URL -- a Zarr store, "file:///data/t.zarr#mode=zarr,file" -- is
+          # not a path.  File.exist? cannot see it, NCZarr does not honour
+          # NC_NOCLOBBER, and a tmp name appended to it lands in the "#mode="
+          # part, so libnetcdf creates over the target itself.
+          if url?(path) && (atomic || clobber != true)
+            raise ArgumentError, "atomic: and clobber: take a file path, not a URL"
+          end
+          # clobber: :skip — path already exists, quietly do nothing.
+          if clobber == :skip && ::File.exist?(path)
+            return nil
+          end
+          # Format inheritance from the template unless netcdf4: is set
+          # explicitly.  A NetCDF-4 template silently upgrades the output
+          # so vars using groups / chunking / deflate stay defineable.
+          if template && !netcdf4
+            src_file = template.is_a?(File) ? template : template.file
+            # nc_inq_format returns the NC_FORMAT_* Integer constant.
+            fmt = src_file.format
+            netcdf4 = true if fmt == ::NetCDF::NC_FORMAT_NETCDF4 ||
+                              fmt == ::NetCDF::NC_FORMAT_NETCDF4_CLASSIC
+          end
+          # Pre-flight clobber check for atomic: refuse before creating a tmp.
+          if atomic && clobber == false && ::File.exist?(path)
+            raise Errno::EEXIST, path
+          end
+          flag  = clobber ? ::NetCDF::NC_CLOBBER : ::NetCDF::NC_NOCLOBBER
+          flag |= ::NetCDF::NC_NETCDF4 if netcdf4
+          write_path = atomic ? atomic_tmp_path_for(path) : path
+          after_new  = template ?
+                       ->(f) { apply_template(f, template, dims: dims, atts: atts, vars: vars) } :
+                       nil
+          build(::NetCDF.nc_create(write_path, flag), path,
+                mode: :create, codec: codec, defining: true,
+                tmp_path: (write_path if atomic),
+                var_defaults: var_defaults,
+                after_new: after_new, &block)
+        end
+
+        private
+
+        # libnetcdf takes a "scheme://..." string as a URL, not a path.
+        def url?(path)
+          path.to_s.match?(%r{\A[A-Za-z][A-Za-z0-9+.-]*://})
+        end
+
+        # Same directory as the target so `File.rename` is same-device
+        # (POSIX cross-device rename fails with EXDEV). Visible (no
+        # dot-prefix) with `.tmp` as the trailing extension so
+        # `find … -name '*.tmp'` picks up orphans. PID + random so
+        # parallel writers don't collide.
+        def atomic_tmp_path_for(path)
+          dir  = ::File.dirname(path)
+          base = ::File.basename(path)
+          ::File.join(dir, "#{base}.#{Process.pid}.#{rand(2**32).to_s(16)}.tmp")
+        end
+
+        def build(ncid, filename, mode:, codec:, defining:,
+                  tmp_path: nil, var_defaults: {}, after_new: nil, &block)
+          f = new(ncid, filename, mode, codec, defining, tmp_path, var_defaults)
+          # after_new runs BEFORE the user's block, but still inside the same
+          # try/close so a failure there also aborts an atomic write.
+          begin
+            after_new&.call(f)
+          rescue Exception
+            f.abort rescue nil
+            raise
+          end
+          return f unless block
+          begin
+            result = block.call(f)
+            f.close                  # normal close: atomic rename happens here
+            result
+          rescue Exception
+            f.abort rescue nil       # abort: unlink tmp, no rename
+            raise
+          end
+        end
+
+        # Copy dims / atts / whole vars from a template File or Group into
+        # a freshly-created File. Called from create() via after_new so
+        # the copy happens BEFORE the user's block runs — the block sees
+        # a File already populated with the template's skeleton and can
+        # start adding new definitions immediately.
+        def apply_template(f, template, dims:, atts:, vars:)
+          case dims
+          when true
+            template.dims.each_value { |d| f.copy_dim(d) }
+          when false, nil
+            # skip
+          when Array
+            dims.each { |n| f.copy_dim(n.to_s, from: template) }
+          else
+            raise ArgumentError,
+                  "File.create dims: expected true / false / Array, got #{dims.inspect}"
+          end
+          case atts
+          when true
+            template.atts.each { |k, v| f.def_att(k, v) }
+          when false, nil
+            # skip
+          when Hash
+            template.atts
+                    .merge(atts.transform_keys(&:to_s))
+                    .each { |k, v| f.def_att(k, v) }
+          else
+            raise ArgumentError,
+                  "File.create atts: expected true / false / Hash, got #{atts.inspect}"
+          end
+          Array(vars).each { |n| f.copy_var(n.to_s, from: template) }
+        end
+        private :apply_template
+      end
+
+      # @api private
+      # Use {.open} / {.create} instead.
+      def initialize(ncid, filename, mode, codec, defining,
+                     tmp_path = nil, var_defaults = {})
+        @filename      = filename
+        @mode          = mode
+        @defining      = defining
+        @codec_request = codec
+        @tmp_path      = tmp_path
+        @var_defaults  = var_defaults.freeze
+        # We are the root: pass self as `file`, nil as `parent`.
+        super(ncid, self, nil)
+        # In create mode the Conventions attribute is not defined yet, so
+        # resolve the codec lazily (on first read/write, once definitions are
+        # in place).  In read mode it can be resolved eagerly.
+        @codec = defining ? nil : Codec.resolve(self, codec)
+        # Safety net: if the user forgets close on an atomic-write File,
+        # at least unlink the tmp (never publish half-written data via a
+        # finalizer — rename is close's job, not GC's).
+        if @tmp_path
+          ObjectSpace.define_finalizer(self,
+            self.class.send(:atomic_tmp_cleaner, @tmp_path))
+        end
+      end
+
+      # Registers the finalizer as a separate class method so the proc
+      # does NOT capture `self` (that would pin the File and defeat GC).
+      # @api private
+      def self.atomic_tmp_cleaner(tmp)
+        proc { ::File.unlink(tmp) if ::File.exist?(tmp) rescue nil }
+      end
+      private_class_method :atomic_tmp_cleaner
+
+      # The final (logical) file path. For an atomic-write File before
+      # close, the bytes are actually at a same-dir tmp path (see
+      # {#tmp_path}); once close succeeds, they are at {#filename}.
+      # (Not to be confused with {Group#path}, which is the NetCDF-4
+      # group hierarchy path — always `"/"` for a root group == a File.)
+      # @return [String]
+      attr_reader :filename
+
+      # The tmp path an atomic-write File is currently writing to; `nil`
+      # otherwise or after a successful close.
+      # @return [String, nil]
+      attr_reader :tmp_path
+
+      # File-wide NetCDF-4 var defaults applied to every `def_var` unless
+      # overridden per-var.
+      # @return [Hash]
+      attr_reader :var_defaults
+
+      # @return [Boolean]
+      def atomic? = !@tmp_path.nil?
+
+      # @return [Object] the codec (resolved lazily on first access)
+      def codec = (@codec ||= Codec.resolve(self, @codec_request))
+
+      # Define a global attribute. Under `codec: :auto`, (re)defining
+      # `Conventions` drops the codec resolved so far, so the next read or
+      # write picks it by the new value — even if something was already
+      # read or written before the attribute was set.
+      # @param (see Group#def_att)
+      # @return [Object] the value written
+      def def_att(name, value)
+        super.tap { @codec = nil if @codec_request == :auto && name.to_s == "Conventions" }
+      end
+
+      # @return [Integer] `NC_FORMAT_*` constant
+      def format = nc_inq_format(@ncid)
+
+      # Close and, for an atomic-write File, atomically rename the tmp to
+      # the final path. Idempotent: a second call is a no-op.
+      # @return [void]
+      def close
+        return if @closed
+        nc_close(@ncid)
+        @closed = true
+        if @tmp_path
+          ::File.rename(@tmp_path, @filename)
+          @tmp_path = nil
+        end
+      end
+
+      # Discard a partially-written atomic file. Closes without renaming
+      # and unlinks the tmp. Used internally by the block-form `ensure`
+      # clause when the block raises; can be called explicitly to abandon
+      # a write in progress. Idempotent.
+      # @return [void]
+      def abort
+        return if @closed
+        nc_close(@ncid) rescue nil
+        @closed = true
+        if @tmp_path
+          ::File.unlink(@tmp_path) if ::File.exist?(@tmp_path)
+          @tmp_path = nil
+        end
+      end
+
+      # @return [Boolean]
+      def closed? = @closed == true
+
+      # The `Conventions` global attribute split into tokens (empty if
+      # the attribute is missing). Provided so a higher-level convenience
+      # layer can dispatch on the same string this gem uses to pick a
+      # codec, without re-parsing it. Values like `"CF-1.7 ACDD-1.3"` or
+      # `"COARDS/CF-1.0"` decompose the way you'd expect.
+      # @return [Array<String>]
+      def conventions
+        s = att("Conventions") or return []
+        s.split(/[\s,\/]+/).reject(&:empty?)
+      end
+
+      # @!group Define / data mode management (also used by every Group / Var)
+
+      # Switch to define mode if not already in it.
+      # @return [void]
+      def ensure_define_mode
+        return if @defining
+        nc_redef(@ncid)
+        @defining = true
+      end
+
+      # Switch to data mode if not already in it.
+      # @return [void]
+      def ensure_data_mode
+        return unless @defining
+        nc_enddef(@ncid)
+        @defining = false
+      end
+
+      # @api private
+      # Run a block within a define-mode (or data-mode) transaction.
+      # @param define [Boolean] `true` to enter define mode first,
+      #   `false` for data mode
+      # @yield the transactional body
+      # @return [Object] the block's return value
+      def __transaction__(define: true)
+        define ? ensure_define_mode : ensure_data_mode
+        yield
+      end
+
+      # @!endgroup
+
+      # @return [String]
+      def inspect = "#<#{self.class} #{@filename.inspect} vars=#{vars.keys} groups=#{groups.keys}>"
+    end
+  end
+end