Merge pull request #648 from AayushSabharwal/as/fntype

feat: add generic to `FnType` representing type of the function
JuliaSymbolics · Sep 16, 2024 · cd7fdb0 · cd7fdb0
2 parents d143b79 + 1dd611e
commit cd7fdb0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/src/types.jl b/src/types.jl
@@ -912,7 +912,7 @@ promote_symtype(f, Ts...) = Any
 #### Function-like variables
 #---------------------------
 
-struct FnType{X<:Tuple,Y} end
+struct FnType{X<:Tuple,Y,Z} end
 
 (f::Symbolic{<:FnType})(args...) = Term{promote_symtype(f, symtype.(args)...)}(f, Any[args...])
 
@@ -1021,8 +1021,16 @@ function _name_type(x)
         lhs, rhs = x.args[1:2]
         if lhs isa Expr && lhs.head === :call
             # e.g. f(::Real)::Unreal
+            if lhs.args[1] isa Expr
+                func_name_and_type = _name_type(lhs.args[1])
+                name = func_name_and_type.name
+                functype = func_name_and_type.type
+            else
+                name = lhs.args[1]
+                functype = Nothing
+            end
             type = map(x->_name_type(x).type, lhs.args[2:end])
-            return (name=lhs.args[1], type=:($FnType{Tuple{$(type...)}, $rhs}))
+            return (name=name, type=:($FnType{Tuple{$(type...)}, $rhs, $functype}))
         else
             return (name=lhs, type=rhs)
         end

diff --git a/test/basics.jl b/test/basics.jl
@@ -6,7 +6,7 @@ using Test
 
 @testset "@syms" begin
     let
-        @syms a b::Float64 f(::Real) g(p, h(q::Real))::Int
+        @syms a b::Float64 f(::Real) g(p, h(q::Real))::Int 
 
         @test issym(a) && symtype(a) == Number
         @test a.name === :a
@@ -16,9 +16,11 @@ using Test
 
         @test issym(f)
         @test f.name === :f
+        @test symtype(f) == FnType{Tuple{Real}, Number, Nothing}
 
         @test issym(g)
         @test g.name === :g
+        @test symtype(g) == FnType{Tuple{Number, FnType{Tuple{Real}, Number, Nothing}}, Int, Nothing}
 
         @test isterm(f(b))
         @test symtype(f(b)) === Number
@@ -32,6 +34,21 @@ using Test
         # issue #91
         @syms h(a,b,c)
         @test isequal(h(1,2,3), h(1,2,3))
+
+        @syms (f::typeof(max))(::Real, ::AbstractFloat)::Number a::Real
+        @test issym(f)
+        @test f.name == :f
+        @test symtype(f) == FnType{Tuple{Real, AbstractFloat}, Number, typeof(max)}
+        @test isterm(f(a, b))
+        @test symtype(f(a, b)) == Number
+
+        @syms g(p, (h::typeof(identity))(q::Real)::Number)::Number
+        @test issym(g)
+        @test g.name == :g
+        @test symtype(g) == FnType{Tuple{Number, FnType{Tuple{Real}, Number, typeof(identity)}}, Number, Nothing}
+        @test_throws "not a subtype of" g(a, f)
+        @syms (f::typeof(identity))(::Real)::Number
+        @test symtype(g(a, f)) == Number
     end
 end