Split checks into per method functions

Author: oxinabox

src/global_checks.jl

@@ -7,6 +7,53 @@ _parameters(sig::UnionAll) = _parameters(sig.body)
 _parameters(sig::DataType) = sig.parameters
 _parameters(sig::Union) = Base.uniontypes(sig)
 
+
+"""
+    test_method_signature(frule|rrule, method)
+
+Tests that the method signature is sensible.
+Right now this just means checking the rule is not being applied to `DataType`, `Union`, or
+`UnionAll`.
+which is easy to do accidentally when writing rules for constructors.
+It happens if you write e.g. `rrule(::typeof(Foo), x)` rather than `rrule(::Type{<:Foo}, x)`.
+This would then actually define `rrule(::DataType, x)`. (or `UnionAll` if `Foo`
+was parametric, or `Union` if `Foo` was a type alias for a `Union`)
+"""
+function test_method_signature end
+
+function test_method_signature(::typeof(rrule), method::Method)
+    @testset "Sensible Constructors" begin
+        function_type = if method.sig <: Tuple{Any, RuleConfig, Type, Vararg}
+            _parameters(method.sig)[3]
+        elseif method.sig <: Tuple{Any, Type, Vararg}
+            _parameters(method.sig)else
+            nothing
+        end
+        
+        @test_msg(
+            "Bad constructor rrule. `typeof(T)` used rather than `Type{T}`. $method",
+            function_type ∉ (DataType, UnionAll, Union)
+        )
+    end
+end
+
+function test_method_signature(::typeof(frule), method::Method)
+    @testset "Sensible Constructors" begin
+        function_type = if method.sig <: Tuple{Any, RuleConfig, Any, Type, Vararg}
+            _parameters(method.sig)[4]
+        elseif method.sig <: Tuple{Any, Any, Type, Vararg}
+            _parameters(method.sig)[3]
+        else
+            nothing
+        end
+
+        @test_msg(
+            "Bad constructor frule. `typeof(T)` used rather than `Type{T}`. $method",
+            function_type ∉ (DataType, UnionAll, Union)
+        )
+    end
+end
+
 """
     test_method_tables()
 
@@ -24,33 +71,11 @@ function test_method_tables()
         # `rrule(::Type{<:Foo}, x)` then that would actually define `rrule(::DataType, x)`
         # which would be bad. This test checks for that and fails if such a method exists.
         for method in methods(rrule)
-            function_type = if method.sig <: Tuple{Any, RuleConfig, Type, Vararg}
-                _parameters(method.sig)[3]
-            elseif method.sig <: Tuple{Any, Type, Vararg}
-                _parameters(method.sig)[2]
-            else
-                nothing
-            end
-            
-            @test_msg(
-                "Bad constructor rrule. `typeof(T)` used rather than `Type{T}`. $method",
-                function_type ∉ (DataType, UnionAll, Union)
-            )
+            test_method_signature(method)
         end
         # frule
         for method in methods(frule)
-            function_type = if method.sig <: Tuple{Any, RuleConfig, Any, Type, Vararg}
-                _parameters(method.sig)[4]
-            elseif method.sig <: Tuple{Any, Any, Type, Vararg}
-                _parameters(method.sig)[3]
-            else
-                nothing
-            end
-
-            @test_msg(
-                "Bad constructor frule. `typeof(T)` used rather than `Type{T}`. $method",
-                function_type ∉ (DataType, UnionAll, Union)
-            )
+            test_method_signature(frule, method)
         end
     end
 end