|
| 1 | +using MacroTools |
| 2 | + |
1 | 3 | """ |
2 | | - default_kwargs(f, [obj = Any]) |
| 4 | + default_kwargs(f::Function, args...; kwargs...) |
3 | 5 |
|
4 | | -Return the default keyword arguments for the function `f`. These defaults may be |
5 | | -derived from the contents or type of the second arugment `obj`. |
| 6 | +Returns a set of default keyword arguments, as a `NamedTuple`, for the function `f` |
| 7 | +depending on an arbitrary number of positional arguments. Any number of these default |
| 8 | +keyword arguments can optionally be overwritten by passing the the keyword as a |
| 9 | +keyword argument to this function. |
| 10 | +""" |
| 11 | +function default_kwargs(f::Function, args...; kwargs...) |
| 12 | + return default_kwargs(f, map(typeof, args)...; kwargs...) |
| 13 | +end |
| 14 | +default_kwargs(f::Function, ::Vararg{<:Type}; kwargs...) = (; kwargs...) |
6 | 15 |
|
7 | | -## Interface |
| 16 | +""" |
| 17 | + @default_kwargs |
8 | 18 |
|
9 | | -Given a function `f`, one can optionally set the default keyword arguments for this |
10 | | -function by specializing either of the following two-argument methods: |
| 19 | +Automatically define a `default_kwargs` method for a given function. This macro should |
| 20 | +be applied before a function definition: |
11 | 21 | ``` |
12 | | -ITensorNetworks.default_kwargs(::typeof(f), prob::AbstractProblem) |
13 | | -ITensorNetworks.default_kwargs(::typeof(f), ::Type{<:AbstractProblem}) |
| 22 | +@default_kwargs astypes = true function f(args...; kwargs...) |
| 23 | + ... |
| 24 | +end |
| 25 | +``` |
| 26 | +If `astypes = true` then the `default_kwargs` method is defined in the |
| 27 | +type domain with respect to `args`, i.e. |
| 28 | +``` |
| 29 | +default_kwargs(::typeof(f), arg::T; kwargs...) # astypes = false |
| 30 | +default_kwargs(::typeof(f), arg::Type{<:T}; kwargs...) # astypes = true |
14 | 31 | ``` |
15 | | -If one does not require the contents of `prob::Prob` to generate the defaults then it is |
16 | | -recommended to dispatch on `Type{<:Prob}` directly (second method) so the defaults |
17 | | -can be accessed without constructing an instance of a `Prob`. |
18 | | -
|
19 | | -The return value of `default_kwargs` should be a `NamedTuple`, and will overwrite any |
20 | | -default values set in the function signature. |
21 | 32 | """ |
22 | | -default_kwargs(f) = default_kwargs(f, Any) |
23 | | -default_kwargs(f, obj) = _default_kwargs_fallback(f, obj) |
24 | | - |
25 | | -# To avoid annoying potential method ambiguities. |
26 | | -function _default_kwargs_fallback(f, iter::RegionIterator) |
27 | | - return default_kwargs(f, problem(iter)) |
28 | | -end |
29 | | -function _default_kwargs_fallback(f, problem::AbstractProblem) |
30 | | - return default_kwargs(f, typeof(problem)) |
| 33 | +macro default_kwargs(args...) |
| 34 | + kwargs = (;) |
| 35 | + for opt in args |
| 36 | + if @capture(opt, key_ = val_) |
| 37 | + @info "" key val |
| 38 | + kwargs = merge(kwargs, NamedTuple{(key,)}((val,))) |
| 39 | + elseif opt === last(args) |
| 40 | + return default_kwargs_macro(opt; kwargs...) |
| 41 | + else |
| 42 | + throw(ArgumentError("Unknown expression object")) |
| 43 | + end |
| 44 | + end |
31 | 45 | end |
32 | 46 |
|
33 | | -# Eventually we reach this if nothing is specialized. |
34 | | -_default_kwargs_fallback(::Any, ::DataType) = (;) |
| 47 | +function default_kwargs_macro(function_def; astypes=true) |
| 48 | + if !isdef(function_def) |
| 49 | + throw( |
| 50 | + ArgumentError("The @default_kwargs macro must be followed by a function definition") |
| 51 | + ) |
| 52 | + end |
35 | 53 |
|
36 | | -""" |
37 | | - current_kwargs(f, iter::RegionIterator) |
| 54 | + ex = splitdef(function_def) |
| 55 | + new_ex = deepcopy(ex) |
38 | 56 |
|
39 | | -Return the keyword arguments to be passed to the function `f` for the current region |
40 | | -defined by the stateful iterator `iter`. |
41 | | -""" |
42 | | -function current_kwargs(f::Function, iter::RegionIterator) |
43 | | - region_kwargs = get(current_region_kwargs(iter), Symbol(f, :_kwargs), (;)) |
44 | | - rv = merge(default_kwargs(f, iter), region_kwargs) |
45 | | - return rv |
46 | | -end |
| 57 | + prev_kwargs = [] |
| 58 | + |
| 59 | + # Give very positional argument a name and escape the type. |
| 60 | + ex[:args] = map(ex[:args]) do arg |
| 61 | + @capture(arg, (name_::T_) | (::T_) | name_) |
| 62 | + if isnothing(name) |
| 63 | + name = gensym() |
| 64 | + end |
| 65 | + if isnothing(T) |
| 66 | + T = :Any |
| 67 | + end |
| 68 | + return :($(name)::$(esc(T))) |
| 69 | + end |
| 70 | + |
| 71 | + # Replacing the kwargs values with the output of `default_kwargs` |
| 72 | + ex[:kwargs] = map(ex[:kwargs]) do kw |
| 73 | + @capture(kw, (key_::T_ = val_) | (key_ = val_) | key_) |
| 74 | + if !isnothing(val) |
| 75 | + kw.args[2] = |
| 76 | + :(default_kwargs($(esc(ex[:name])), $(ex[:args]...); $(prev_kwargs...)).$key) |
| 77 | + end |
| 78 | + push!(prev_kwargs, key) |
| 79 | + return kw |
| 80 | + end |
| 81 | + |
| 82 | + # Promote to the type domain if wanted |
| 83 | + if astypes |
| 84 | + new_ex[:args] = map(ex[:args]) do arg |
| 85 | + @capture(arg, name_::T_) |
| 86 | + return :($(name)::Type{<:$T}) |
| 87 | + end |
| 88 | + end |
| 89 | + |
| 90 | + new_ex[:name] = :(ITensorNetworks.default_kwargs) |
| 91 | + new_ex[:args] = convert(Vector{Any}, ex[:args]) |
47 | 92 |
|
48 | | -# Generic |
| 93 | + new_ex[:args] = pushfirst!(new_ex[:args], :(::typeof($(esc(ex[:name]))))) |
49 | 94 |
|
50 | | -# I think these should be set independent of a function, but for now: |
51 | | -function default_kwargs(::typeof(factorize), ::Any) |
52 | | - return (; maxdim=typemax(Int), cutoff=0.0, mindim=1) |
| 95 | + # Escape anything on the right-hand side of a keyword definition. |
| 96 | + new_ex[:kwargs] = map(new_ex[:kwargs]) do kw |
| 97 | + @capture(kw, (key_ = val_) | key_) |
| 98 | + if !isnothing(val) |
| 99 | + kw.args[2] = esc(val) |
| 100 | + end |
| 101 | + return kw |
| 102 | + end |
| 103 | + |
| 104 | + new_ex[:body] = :(return (; $(prev_kwargs...))) |
| 105 | + |
| 106 | + # Escape the actual function name |
| 107 | + ex[:name] = :($(esc(ex[:name]))) |
| 108 | + |
| 109 | + rv = quote |
| 110 | + $(combinedef(ex)) |
| 111 | + $(combinedef(new_ex)) |
| 112 | + end |
| 113 | + |
| 114 | + return rv |
53 | 115 | end |
0 commit comments