Added save compatible version of filenames, arrays and images function

src/collection.jl

@@ -9,9 +9,66 @@ end
 
 parse_name(filename, ext, ::Val{true}) = filename
 
+
+# utility function for generating filename
+function generate_filename(path, save_location, save_prefix, save_suffix, save_delim, ext)
+    # get the filename
+    filename = last(splitdir(path))
+
+    # splitting name and extension
+    modified_name, ext = parse_name_ext(filename, "." * ext)
+
+    # adding prefix and suffix with delimiter
+    if !isnothing(save_prefix)
+        modified_name = string(save_prefix, save_delim, modified_name)
+    end
+    if !isnothing(save_suffix)
+        modified_name = string(modified_name, save_delim, save_suffix)
+    end
+
+    # adding extension to modified_name
+    if ext == ""
+        file_path = joinpath(save_location, modified_name * ".fits")
+    else
+        file_path = joinpath(save_location, modified_name * ext)
+    end
+    return file_path
+end
+
+
+# utility function to return filename and extension separately
+# returns extension including "." at the beginning
+function parse_name_ext(filename, ext)
+    idxs = findall(ext, filename)
+    length(idxs) == 0 && return (filename, "")
+    breaking_index = first(last(idxs))
+    return filename[1:breaking_index - 1], filename[breaking_index:end]
+end
+
+
+"""
+    write_data(file_path, data)
+
+Writes `data` in FITS format at `file_path`.
+"""
+function write_data(file_path, data)
+    d = ndims(data)
+    transposed_data = permutedims(data, d:-1:1)
+    FITS(file_path, "w") do fh
+        write(fh, transposed_data)
+    end
+end
+
 #---------------------------------------------------------------------------------------
 @doc raw"""
-    fitscollection(dir; recursive=true, abspath=true, keepext=true, ext=r"fits(\.tar\.gz)?", exclude=nothing, exclude_dir=nothing, exclude_key = ("", "HISTORY"))
+    fitscollection(dir;
+                   recursive=true,
+                   abspath=true,
+                   keepext=true,
+                   ext=r"fits(\.tar\.gz)?",
+                   exclude=nothing,
+                   exclude_dir=nothing,
+                   exclude_key=("", "HISTORY"))
 
 Walk through `dir` collecting FITS files, scanning their headers, and culminating into a `DataFrame` that can be used with the generators for iterating over many files and processing them. If `recursive` is false, no subdirectories will be walked through.
 
@@ -60,7 +117,7 @@ function fitscollection(basedir::String;
                         exclude = nothing,
                         exclude_dir = nothing,
                         exclude_key = ("", "HISTORY"))
-    df = DataFrame()
+    collection = DataFrame()
 
     for (root, dirs, files) in walkdir(basedir)
         # recursive searching functionality
@@ -88,25 +145,25 @@ function fitscollection(basedir::String;
                 # filtering out comment columns
                 _keys = filter(k -> k ∉ exclude_key, keys(header_data))
                 _values = (header_data[k] for k in _keys)
-                push!(df, (path = path, name = name, hdu = index, zip(Symbol.(_keys), _values)...); cols = :union)
+                push!(collection, (path = path, name = name, hdu = index, zip(Symbol.(_keys), _values)...); cols = :union)
             end
             close(fits_data)
         end
     end
-    return df
+    return collection
 end
 
 
 """
-    arrays(df::DataFrame)
+    arrays(collection::DataFrame)
 
 Generator for arrays of images of entries in data frame.
 """
 function arrays end
 
 # generator for image arrays specified by data frames (i.e. path of file, hdu etc.)
-@resumable function arrays(df::DataFrame)
-    for row in eachrow(df)
+@resumable function arrays(collection::DataFrame)
+    for row in eachrow(collection)
         fh = FITS(row.path)
         @yield getdata(fh[row.hdu])
         close(fh)
@@ -115,30 +172,107 @@ end
 
 
 """
-    filenames(df::DataFrame)
+    filenames(collection::DataFrame)
 
 Generator for filenames of entries in data frame.
 """
 function filenames end
 
 # generator for filenames specified by data frame (i.e. path of file, hdu etc.)
-@resumable function filenames(df::DataFrame)
-    for row in eachrow(df)
+@resumable function filenames(collection::DataFrame)
+    for row in eachrow(collection)
         @yield row.path
     end
 end
 
 
 """
-    images(df::DataFrame)
+    images(collection::DataFrame)
 
 Generator for `ImageHDU`s of entries in data frame.
 """
 function images end
 
 # generator for ImageHDU specified by data frame (i.e. path of file, hdu etc.)
-@resumable function images(df::DataFrame)
-    for row in eachrow(df)
+@resumable function images(collection::DataFrame)
+    for row in eachrow(collection)
         @yield FITS(row.path)[row.hdu]
     end
 end
+
+
+"""
+    images(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\\.tar\\.gz)?"i)
+
+Applies function `f` on all ImageHDUs present in data frame and saves it in FITS file.
+
+It returns an array of output values of function `f` applied on ImageHDUs. If `save = false`, the save functionality does not execute. File is saved at `path` specified by the user.
+Suffix and prefix can be added to filename of newly created files by modifying `save_suffix` and `save_prefix`, `save_delim` is used as delimiter.
+`ext` is the extension of files to be taken into consideration for applying function, by default it is set to `r"fits(\\.tar\\.gz)?"i`.
+"""
+function images(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\.tar\.gz)?"i, kwargs...)
+    image_iterator = images(collection; kwargs...)
+    locations = collection.path
+
+    processed_images = map(zip(locations, image_iterator)) do (location, output)
+        processed_image = f(output)
+        if save
+            save_path = generate_filename(location, path, save_prefix, save_suffix, save_delim, ext)
+            write_data(save_path, processed_image)
+        end
+        processed_image
+    end
+
+    return processed_images
+end
+
+"""
+    filenames(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\\.tar\\.gz)?"i)
+
+Applies function `f` on all file paths present in data frame and saves it in FITS file.
+
+It returns an array of output values of function `f` applied on file paths. If `save = false`, the save functionality does not execute. File is saved at `path` specified by the user.
+Suffix and prefix can be added to filename of newly created files by modifying `save_suffix` and `save_prefix`, `save_delim` is used as delimiter.
+`ext` is the extension of files to be taken into consideration for applying function, by default it is set to `r"fits(\\.tar\\.gz)?"i`.
+"""
+function filenames(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\.tar\.gz)?"i, kwargs...)
+    path_iterator = filenames(collection; kwargs...)
+    locations = collection.path
+
+    processed_images = map(zip(locations, path_iterator)) do (location, output)
+        processed_image = f(output)
+        if save
+            save_path = generate_filename(location, path, save_prefix, save_suffix, save_delim, ext)
+            write_data(save_path, processed_image)
+        end
+        processed_image
+    end
+
+    return processed_images
+end
+
+
+"""
+    arrays(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\\.tar\\.gz)?"i)
+
+Applies function `f` on all image arrays present in data frame and saves it in FITS file.
+
+It returns an array of output values of function `f` applied on image arrays. If `save = false`, the save functionality does not execute. File is saved at `path` specified by the user.
+Suffix and prefix can be added to filename of newly created files by modifying `save_suffix` and `save_prefix`, `save_delim` is used as delimiter.
+`ext` is the extension of files to be taken into consideration for applying function, by default it is set to `r"fits(\\.tar\\.gz)?"i`.
+"""
+function arrays(f, collection::DataFrame; save = false, path = nothing, save_prefix = nothing, save_suffix = nothing, save_delim = "_", ext = r"fits(\.tar\.gz)?"i, kwargs...)
+    array_iterator = arrays(collection; kwargs...)
+    locations = collection.path
+
+    processed_images = map(zip(locations, array_iterator)) do (location, output)
+        processed_image = f(output)
+        if save
+            save_path = generate_filename(location, path, save_prefix, save_suffix, save_delim, ext)
+            write_data(save_path, processed_image)
+        end
+        processed_image
+    end
+
+    return processed_images
+end