use LazyModules.jl

CarloLucibello · CarloLucibello · commit 1656650b3fc1 · 2022-05-17T07:49:03.000+02:00
diff --git a/Project.toml b/Project.toml
@@ -16,6 +16,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+LazyModules = "8cdb02fc-e678-4876-92c5-9defec4f444e"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
@@ -37,6 +38,7 @@ HDF5 = "0.16.2"
 ImageCore = "0.9"
 JLD2 = "0.4.21"
 JSON3 = "1"
+LazyModules = "0.3"
 MAT = "0.10"
 MLUtils = "0.2.0"
 NPZ = "0.4.1"
diff --git a/docs/src/containers/overview.md b/docs/src/containers/overview.md
@@ -5,10 +5,15 @@ MLDatasets.jl contains several reusable data containers for accessing datasets i
 ```@docs
 FileDataset
 TableDataset
-HDF5Dataset
-Base.close(::HDF5Dataset)
-JLD2Dataset
-Base.close(::JLD2Dataset)
 CachedDataset
 MLDatasets.make_cache
 ```
+
+<!-- 
+# TODO add back to docs when included again in the pkg
+HDF5Dataset
+Base.close(::HDF5Dataset)
+JLD2Dataset
+Base.close(::JLD2Dataset) 
+-->
+
diff --git a/src/MLDatasets.jl b/src/MLDatasets.jl
@@ -3,33 +3,40 @@ module MLDatasets
 using FixedPointNumbers
 using SparseArrays
 using Tables
-# using Glob
-# using DataFrames
-# import ImageCore
 using DataDeps
 import MLUtils
 using MLUtils: getobs, numobs, AbstractDataContainer
 using ColorTypes
-
-### I/O imports
-# import NPZ
-# import Pickle
-# using MAT: matopen, matread
-using FileIO
-# import CSV 
-# using HDF5
-# using JLD2
-# import JSON3
+using Glob
 using DelimitedFiles: readdlm
-##########
+using FileIO
+using LazyModules: @lazy
+
+include("require.jl") # export @require
+
+# Use `@lazy import SomePkg` whenever the returned types are not its own types,
+# since for methods applied on the returned types we would encounter in world-age issues
+# (see discussion in  https://github.com/JuliaML/MLDatasets.jl/pull/128).
+# In the other case instead, use `require import SomePkg` to force 
+# the use to manually import.
+
+@require import JSON3="0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+@require import DataFrames="a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+@lazy import ImageCore="a09fc81d-aa75-5fe9-8630-4744c3626534"
+# @lazy import NPZ # lazy imported by FileIO
+@lazy import Pickle="fbb45041-c46e-462f-888f-7c521cafbc2c"
+@lazy import MAT="23992714-dd62-5051-b70f-ba57cb901cac"
+@lazy import CSV="336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+@lazy import HDF5="f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+# @lazy import JLD2
 
 export getobs, numobs # From MLUtils.jl
 
 include("abstract_datasets.jl")
 # export AbstractDataset, 
 #        SupervisedDataset
 
-include("imports.jl")
+# include("imports.jl")
 include("utils.jl")
 export convert2image
 
@@ -39,16 +46,16 @@ include("io.jl")
 include("download.jl")
 
 ### API to be revisited with conditional module loading
-# include("containers/filedataset.jl")
-# export FileDataset
-# include("containers/tabledataset.jl")
-# export TableDataset
+include("containers/filedataset.jl")
+export FileDataset
+include("containers/tabledataset.jl")
+export TableDataset
 # include("containers/hdf5dataset.jl")
 # export HDF5Dataset
 # # include("containers/jld2dataset.jl")
 # # export JLD2Dataset
-# include("containers/cacheddataset.jl")
-# export CachedDataset
+include("containers/cacheddataset.jl")
+export CachedDataset
 
 # Misc.
 include("datasets/misc/boston_housing.jl")
diff --git a/src/containers/tabledataset.jl b/src/containers/tabledataset.jl
@@ -50,8 +50,8 @@ function Base.length(dataset::TableDataset)
 end
 
 # fast access for DataFrame
-Base.getindex(dataset::TableDataset{<:DataFrame}, i) = dataset.table[i, :]
-Base.length(dataset::TableDataset{<:DataFrame}) = nrow(dataset.table)
+# Base.getindex(dataset::TableDataset{<:DataFrame}, i) = dataset.table[i, :]
+# Base.length(dataset::TableDataset{<:DataFrame}) = nrow(dataset.table)
 
 # fast access for CSV.File
 # Base.getindex(dataset::TableDataset{<:CSV.File}, i) = dataset.table[i]
diff --git a/src/datasets/misc/boston_housing.jl b/src/datasets/misc/boston_housing.jl
@@ -78,8 +78,7 @@ function BostonHousing(; as_df = true, dir = nothing)
     @assert dir === nothing "custom `dir` is not supported at the moment."
     path = joinpath(@__DIR__, "..", "..", "..", "data", "boston_housing.csv")
     df = read_csv(path)
-    DFs = require_import(:DataFrames)
-    features = df[!, DFs.Not(:MEDV)]
+    features = df[!, DataFrames.Not(:MEDV)]
     targets = df[!, [:MEDV]]
     
     metadata = Dict{String, Any}()
diff --git a/src/datasets/misc/iris.jl b/src/datasets/misc/iris.jl
@@ -80,10 +80,9 @@ end
 function Iris(; dir = nothing, as_df = true)
     path = datafile("Iris", "iris.data", dir)
     df = read_csv(path, header=0)
-    DFs = require_import(:DataFrames)
-    DFs.rename!(df, ["sepallength", "sepalwidth", "petallength", "petalwidth", "class"])
+    DataFrames.rename!(df, ["sepallength", "sepalwidth", "petallength", "petalwidth", "class"])
 
-    features = df[!, DFs.Not(:class)]
+    features = df[!, DataFrames.Not(:class)]
     targets = df[!, [:class]]
 
     metadata = Dict{String, Any}()
diff --git a/src/datasets/misc/titanic.jl b/src/datasets/misc/titanic.jl
@@ -62,9 +62,8 @@ function Titanic(; as_df = true, dir = nothing)
     @assert dir === nothing "custom `dir` is not supported at the moment."
     path = joinpath(@__DIR__, "..", "..", "..", "data", "titanic.csv")
     df = read_csv(path)
-    DFs = require_import(:DataFrames)
-    
-    features = df[!, DFs.Not(:Survived)]
+
+    features = df[!, DataFrames.Not(:Survived)]
     targets = df[!, [:Survived]]
     
     metadata = Dict{String, Any}()
diff --git a/src/datasets/vision/cifar10.jl b/src/datasets/vision/cifar10.jl
@@ -165,7 +165,6 @@ convert2image(::Type{<:CIFAR10}, x::AbstractArray{<:Integer}) =
 function convert2image(::Type{<:CIFAR10}, x::AbstractArray{T,N}) where {T,N}
     @assert N == 3 || N == 4
     x = permutedims(x, (3, 2, 1, 4:N...))
-    ImageCore = lazy_import(:ImageCore)
     return ImageCore.colorview(RGB, x)
 end
 
diff --git a/src/datasets/vision/mnist.jl b/src/datasets/vision/mnist.jl
@@ -141,7 +141,6 @@ convert2image(::Type{<:MNIST}, x::AbstractArray{<:Integer}) =
 function convert2image(::Type{<:MNIST}, x::AbstractArray{T,N}) where {T,N}
     @assert N == 2 || N == 3
     x = permutedims(x, (2, 1, 3:N...))
-    ImageCore = lazy_import(:ImageCore)
     return ImageCore.colorview(Gray, x)
 end
 
diff --git a/src/imports.jl b/src/imports.jl
diff --git a/src/io.jl b/src/io.jl
@@ -1,7 +1,3 @@
-# IMPORTS 
-# Use `lazy_import(:SomePkg)` whenever the returned types are "native" types,
-#  i.e. not defined in SomePkg itself, otherwise use `require_import(:SomePkg)`.
-
 
 function read_csv(path; kws...)
     return read_csv_asdf(path; kws...)
@@ -17,8 +13,6 @@ function read_csv(path, sink::Type{A}; kws...) where A <: AbstractMatrix
 end
 
 function read_csv_asdf(path; kws...)
-    DataFrames = require_import(:DataFrames)
-    CSV = lazy_import(:CSV)
     return CSV.read(path, DataFrames.DataFrame; kws...)
 end
 
@@ -31,21 +25,17 @@ function read_npz(path)
 end
 
 function read_pytorch(path)
-    Pickle = lazy_import(:Pickle)
     return Pickle.Torch.THload(path)
 end
 
 function read_pickle(path)
-    Pickle = lazy_import(:Pickle)
     return Pickle.npyload(path)
 end
 
 function read_mat(path)
-    MAT = lazy_import(:MAT)
     return MAT.matread(path)
 end
 
 function read_json(path)
-    JSON3 = require_import(:JSON3)
     return open(JSON3.read, path)
 end
diff --git a/src/require.jl b/src/require.jl
@@ -0,0 +1,121 @@
+
+using Base: invokelatest
+using Base: PkgId, UUID
+
+# export @require
+
+const _LOAD_LOCKER = Threads.ReentrantLock()
+
+mutable struct RequireModule
+    _require_pkgid::PkgId
+    _require_loaded::Bool
+end
+RequireModule(uuid::UUID, name::String) = RequireModule(PkgId(uuid, name), false)
+Base.Docs.Binding(m::RequireModule, v::Symbol) = Base.Docs.Binding(checked_import(m._require_pkgid), v)
+function Base.show(io::IO, m::RequireModule)
+    print(io, "RequireModule(", m._require_pkgid.name, ")")
+end
+
+function Base.getproperty(m::RequireModule, s::Symbol)
+    if s in (:_require_pkgid, :_require_loaded)
+        return getfield(m, s)
+    end
+    lm = require_import(m)
+    return getproperty(lm, s)
+end
+
+function assert_imported(pkgid::PkgId)
+    if !Base.root_module_exists(pkgid)
+        name = pkgid.name
+        error("Add `import $name` or `using $name` to your code to unlock this functionality.")
+    end
+end
+
+function require_import(m::RequireModule)
+    pkgid = getfield(m, :_require_pkgid)
+    if !getfield(m, :_require_loaded)
+        assert_imported(pkgid)
+        setfield!(m, :_require_loaded, true)
+    end
+    return Base.root_module(pkgid)
+end
+
+
+"""
+    @require import PkgName=UUID
+
+Force the user to add to their code `import package PkgName`.
+only if they want the functionality provided by PkgName, 
+otherwise PkgName won't be imported and won't affect loading time.
+
+```julia
+module MyRequirePkg
+    @require import Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+    
+    draw_figure(data) = Plots.plot(data, title="MyPkg Plot")
+end
+```
+"""
+macro require(ex)
+    if ex.head in (:import, :using)
+        error("require `@require $(ex)=UUID` format.")
+    end
+    if ex.head != :(=)
+        error("unrecognized expression: $(ex)")
+    end
+    uuid = UUID(ex.args[2])
+    ex = ex.args[1]
+
+    if ex.head != :import
+        @warn "only `import` command is supported, fallback to eager mode"
+        return ex
+    end
+    args = ex.args
+    if length(args) != 1
+        @warn "only single package import is supported, fallback to eager mode"
+        return ex
+    end
+    x = args[1]
+    if x.head == :.
+        # usage: @require import Foo
+        m = _require_load(__module__, x.args[1], uuid, x.args[1])
+        # TODO(johnnychen94): the background eager loading seems to work only for Main scope
+        isa(m, Module) && return m
+        isnothing(m) && return ex
+        return m
+    elseif x.head == :(:)
+        # usage: @require import Foo: foo, bar
+        @warn "lazily importing symbols are not supported, fallback to eager mode"
+        return ex
+    elseif x.head == :as # compat: Julia at least v1.6
+        # usage: @require import Foo as RequireFoo
+        m = _require_load(__module__, x.args[2], uuid, x.args[1].args[1])
+        isa(m, Module) && return m
+        isnothing(m) && return ex
+        return m
+    else
+        @warn "unrecognized syntax $ex"
+        return ex
+    end
+end
+
+function _require_load(mod, name::Symbol, uuid::UUID, sym::Symbol)
+    if isdefined(mod, name)
+        # otherwise, Revise will constantly trigger the constant redefinition warning
+        m = getfield(mod, name)
+        if m isa RequireModule || m isa Module
+            return m
+        else
+            @warn "Failed to import module, the name `$name` already exists, do nothing"
+            return nothing
+        end
+    end
+    try
+        m = RequireModule(uuid, String(sym))
+        Core.eval(mod, :(const $(name) = $m))
+        return m
+    catch err
+        @warn err
+        return nothing
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -5,8 +5,7 @@ using MLDatasets: Graph
 using DataFrames, Tables
 using ColorTypes
 using FixedPointNumbers
-
-using Pickle, JSON3
+using JSON3
 
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 
