Skip to content

Commit dba3649

Browse files
committed
Split checks into per method functions
1 parent 9045ce2 commit dba3649

File tree

1 file changed

+49
-24
lines changed

1 file changed

+49
-24
lines changed

src/global_checks.jl

Lines changed: 49 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,53 @@ _parameters(sig::UnionAll) = _parameters(sig.body)
77
_parameters(sig::DataType) = sig.parameters
88
_parameters(sig::Union) = Base.uniontypes(sig)
99

10+
11+
"""
12+
test_method_signature(frule|rrule, method)
13+
14+
Tests that the method signature is sensible.
15+
Right now this just means checking the rule is not being applied to `DataType`, `Union`, or
16+
`UnionAll`.
17+
which is easy to do accidentally when writing rules for constructors.
18+
It happens if you write e.g. `rrule(::typeof(Foo), x)` rather than `rrule(::Type{<:Foo}, x)`.
19+
This would then actually define `rrule(::DataType, x)`. (or `UnionAll` if `Foo`
20+
was parametric, or `Union` if `Foo` was a type alias for a `Union`)
21+
"""
22+
function test_method_signature end
23+
24+
function test_method_signature(::typeof(rrule), method::Method)
25+
@testset "Sensible Constructors" begin
26+
function_type = if method.sig <: Tuple{Any, RuleConfig, Type, Vararg}
27+
_parameters(method.sig)[3]
28+
elseif method.sig <: Tuple{Any, Type, Vararg}
29+
_parameters(method.sig)else
30+
nothing
31+
end
32+
33+
@test_msg(
34+
"Bad constructor rrule. `typeof(T)` used rather than `Type{T}`. $method",
35+
function_type (DataType, UnionAll, Union)
36+
)
37+
end
38+
end
39+
40+
function test_method_signature(::typeof(frule), method::Method)
41+
@testset "Sensible Constructors" begin
42+
function_type = if method.sig <: Tuple{Any, RuleConfig, Any, Type, Vararg}
43+
_parameters(method.sig)[4]
44+
elseif method.sig <: Tuple{Any, Any, Type, Vararg}
45+
_parameters(method.sig)[3]
46+
else
47+
nothing
48+
end
49+
50+
@test_msg(
51+
"Bad constructor frule. `typeof(T)` used rather than `Type{T}`. $method",
52+
function_type (DataType, UnionAll, Union)
53+
)
54+
end
55+
end
56+
1057
"""
1158
test_method_tables()
1259
@@ -24,33 +71,11 @@ function test_method_tables()
2471
# `rrule(::Type{<:Foo}, x)` then that would actually define `rrule(::DataType, x)`
2572
# which would be bad. This test checks for that and fails if such a method exists.
2673
for method in methods(rrule)
27-
function_type = if method.sig <: Tuple{Any, RuleConfig, Type, Vararg}
28-
_parameters(method.sig)[3]
29-
elseif method.sig <: Tuple{Any, Type, Vararg}
30-
_parameters(method.sig)[2]
31-
else
32-
nothing
33-
end
34-
35-
@test_msg(
36-
"Bad constructor rrule. `typeof(T)` used rather than `Type{T}`. $method",
37-
function_type (DataType, UnionAll, Union)
38-
)
74+
test_method_signature(method)
3975
end
4076
# frule
4177
for method in methods(frule)
42-
function_type = if method.sig <: Tuple{Any, RuleConfig, Any, Type, Vararg}
43-
_parameters(method.sig)[4]
44-
elseif method.sig <: Tuple{Any, Any, Type, Vararg}
45-
_parameters(method.sig)[3]
46-
else
47-
nothing
48-
end
49-
50-
@test_msg(
51-
"Bad constructor frule. `typeof(T)` used rather than `Type{T}`. $method",
52-
function_type (DataType, UnionAll, Union)
53-
)
78+
test_method_signature(frule, method)
5479
end
5580
end
5681
end

0 commit comments

Comments
 (0)