add collection of all type annotations, not only method sigs

src/types-analysis/data.jl

@@ -14,9 +14,9 @@ JlASTTypeExpr = Any
 Kind of a type annotation:
 - complete method type signature (without return type)
 - return type
-- type assertion (in the method body)
+- type assertion (in the method body) or type annotation (in the field)
 """
-@enum TypeAnnKind mtsig retty tyass
+@enum TypeAnnKind mtsig retty tyassorann
 
 "Information about a type annotation in some file"
 struct TypeAnnInfo
@@ -31,6 +31,8 @@ TypeAnnInfoList = LinkedList{TypeAnnInfo}
 "Data returned by `MacroTools.splitdef`"
 SplitFunDef = Dict{Symbol, Any}
 
+NOTAFUNSIG = "<NOT A FUNSIG>"
+
 #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 # Analyzing type annotations
 #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

src/types-analysis/types-extract.jl

@@ -2,16 +2,16 @@
 # Extracting method type signatures from the source code
 ###############################
 #
-# The goal is to extract all type annotations from method definitions.
+# The goal is to extract all type annotations from method definitions
+# and field declarations.
 #
-# There are 3 sources of type annotations:
+# There are 4 sources of type annotations:
 #
 #   1) method arguments signature (e.g. `f(x::Vector{T}) where T`
 #      corresponds to `Tuple{Vector{T}} where T`)
 #   2) return type annotation (e.g. `f() :: Int`)
 #   3) type assertions in the method body (e.g. `x :: Int`)
-#
-# We start with 1) and ignore 2-3 for now.
+#   4) field type annotations
 #######################################################################
 
 parseAndCollectTypeAnnotations(
@@ -27,38 +27,62 @@ Collects type annotations from all method definitions in `expr`
 """
 collectTypeAnnotations(expr) :: TypeAnnInfoList = begin
     tyAnns = nil(TypeAnnInfo)
-    recordFunDefTyAnns(e) = begin
+    recordTyAnns(e) = begin
         try 
-            tyAnns = collectFunDefTypeAnnotations(e, tyAnns)
+            if isFunDef(e)
+                (e, tyAnns) = collectFunDefTypeSignature(e, tyAnns)
+            end
+            (e, tyAnns) = collectDeclaredTypeAnnotation(e, tyAnns)
         catch err
             @error "Couldn't process expression" e err
         end
-        e # return the same expr for `prewalk` to work
+        e # return expr for `prewalk` to work
     end
-    MacroTools.prewalk(recordFunDefTyAnns, expr)
+    MacroTools.prewalk(recordTyAnns, expr)
     tyAnns
 end
 
 """
-    :: (Any, TypeAnnInfoList) → TypeAnnInfoList
-Returns type annotations from `expr` concatenated with the list `tyAnns`
-if `expr` is a method definition.
+    :: (Any, TypeAnnInfoList) → (Any, TypeAnnInfoList)
+ASSUME: `expr` is a method definition.
 
-FIXME: currently only extracts a method signature,
-but we might be interested in type assertions, too
+Returns:
+- the body of the method encoded by `expr`
+- method signature type and return type (if present) concatenated
+  with the list `tyAnns`.
 """
-collectFunDefTypeAnnotations(expr, tyAnns :: TypeAnnInfoList) = begin
-    isFunDef(expr) || return tyAnns
+collectFunDefTypeSignature(expr, tyAnns :: TypeAnnInfoList) = begin
     funDefParts = splitdef(expr)
     methodArgTuple = getMethodTupleType(funDefParts)
-    cons(
-        TypeAnnInfo(
-            get(funDefParts, :name, "<NA-name>"), 
-            mtsig, 
-            methodArgTuple
-        ),
-        tyAnns
-    )
+    fname = get(funDefParts, :name, "<NA-name>")
+    tyAnns = cons(
+        TypeAnnInfo(fname, mtsig, methodArgTuple),
+        tyAnns)
+    if haskey(funDefParts, :rtype)
+        tyAnns = cons(
+            TypeAnnInfo(fname, retty, funDefParts[:rtype]),
+            tyAnns)
+    end
+    (get(funDefParts, :body, "<NA-body>"), tyAnns)
+end
+
+"""
+    :: (Any, TypeAnnInfoList) → (Any, TypeAnnInfoList)
+ASSUME: `expr` is NOT a method definition
+(nothing bad happens if it is a method definition, but this method is supposed
+to be called for something else to avoid the collection of duplicate annotations)
+
+Returns:
+- `expr` as is or with the type annotation removed if it had one
+- type annotation `ty` if `expr` is `e :: ty` concatenated
+  with the list `tyAnns`, or `tyAnns` otherwise
+"""
+collectDeclaredTypeAnnotation(expr, tyAnns :: TypeAnnInfoList) = begin
+    if @capture(expr, E_ :: TY_)
+        tyAnns = cons(TypeAnnInfo(NOTAFUNSIG, tyassorann, TY), tyAnns)
+        expr = E
+    end
+    (expr, tyAnns)
 end
 
 """

test/types-analysis.jl

@@ -4,9 +4,9 @@
 
 using Main.JuliaSub: TypesAnlsBadMethodParamAST
 
-using Main.JuliaSub: TypeAnnInfo, mtsig, retty, tyass
+using Main.JuliaSub: TypeAnnInfo, mtsig, retty, tyassorann, NOTAFUNSIG
 using Main.JuliaSub: getArgTypeAnn, getMethodTupleType
-using Main.JuliaSub: collectFunDefTypeAnnotations, collectTypeAnnotations
+using Main.JuliaSub: collectFunDefTypeSignature, collectTypeAnnotations
 using Main.JuliaSub: parseAndCollectTypeAnnotations
 
 using Main.JuliaSub: TyVarSummary, TypeTyVarsSummary
@@ -31,6 +31,7 @@ dfedNoArgRet = :(
 
 fdefArgs = :(
     function fdefArgs(x, y::Vector{Bool}, z::Int = 0)
+        baz :: Vector{<:Int}
         "blah"
     end
 )
@@ -57,7 +58,7 @@ fdefNoArgName = :(
 )
 
 fdefQualName = :(
-    Base.fdefQualName(x::Int) = 0
+    Base.fdefQualName(x::Int) :: Int = 0
 )
 
 fdefsSimpleX2 = :(begin
@@ -71,9 +72,13 @@ fdefsSimpleNested = :(module X
 
     struct Bar end
 
+    struct Zoo
+        x :: Bar
+    end
+
     function baz(xs::Vector{T}, y::T) where T
         bar() = Bar()
-        "abc"
+        "abc" :: String
     end
 
     println()
@@ -133,33 +138,41 @@ end
     @test tts[:fdefNoArgName] == :( Tuple{Int} )
 end
 
-# FIXME: collectFunDefTypeAnnotations currently extracts only method signature
-@testset "types-analysis :: collect all ty-anns in mtdef" begin
+@testset "types-analysis :: collect ty sig in mtdef     " begin
     tyAnns = nil(TypeAnnInfo)
+    e = :()
 
-    tyAnns = collectFunDefTypeAnnotations(fdefNoArg, tyAnns)
+    (e, tyAnns) = collectFunDefTypeSignature(fdefNoArg, tyAnns)
     @test tyAnns ==
         list(TypeAnnInfo(:fdefNoArg, mtsig, :( Tuple{} )))
 
-    tyAnns = collectFunDefTypeAnnotations(fdefWhereTriv, tyAnns)
+    (e, tyAnns) = collectFunDefTypeSignature(fdefWhereTriv, tyAnns)
     @test tyAnns ==
         list(
             TypeAnnInfo(:fdefWhereTriv, mtsig, :( Tuple{T} where T )),
             TypeAnnInfo(:fdefNoArg, mtsig, :( Tuple{} ))
         )
 
-    @test collectFunDefTypeAnnotations(fdefWhereRetSimp, nil(TypeAnnInfo)) ==
-        list(TypeAnnInfo(
-            :fdefWhereRetSimp, mtsig, 
-            :( Tuple{Dict{T, S}, Any} where S where T )))
+    @test collectFunDefTypeSignature(fdefWhereRetSimp, nil(TypeAnnInfo))[2] ==
+        list(
+            TypeAnnInfo(
+                :fdefWhereRetSimp, retty,
+                :( Vector{T} ) 
+            ),
+            TypeAnnInfo(
+                :fdefWhereRetSimp, mtsig, 
+                :( Tuple{Dict{T, S}, Any} where S where T )
+            )
+        )
 
-    @test collectFunDefTypeAnnotations(fdefQualName, nil(TypeAnnInfo)) == 
-        list(TypeAnnInfo(:(Base.fdefQualName), mtsig, :( Tuple{Int} )))
-
-    @test collectFunDefTypeAnnotations(:(f()), nil(TypeAnnInfo)) == nil()
+    @test collectFunDefTypeSignature(fdefQualName, nil(TypeAnnInfo))[2] == 
+        list(
+            TypeAnnInfo(:(Base.fdefQualName), retty, :Int),
+            TypeAnnInfo(:(Base.fdefQualName), mtsig, :( Tuple{Int} ))
+        )
 end
 
-# FIXME: collectFunDefTypeAnnotations currently extracts only method signature
+# FIXME: collectFunDefTypeSignature currently extracts only method signature
 @testset "types-analysis :: collect all ty-anns in expr " begin
     @test collectTypeAnnotations(:(x)) == nil()
 
@@ -169,20 +182,28 @@ end
     )
 
     @test collectTypeAnnotations(fdefsSimpleNested) == list(
+        TypeAnnInfo(NOTAFUNSIG, tyassorann, :String),
         TypeAnnInfo(:bar, mtsig, :( Tuple{} )),
         TypeAnnInfo(:baz, mtsig, :( Tuple{Vector{T}, T} where T )),
+        TypeAnnInfo(NOTAFUNSIG, tyassorann, :Bar),
         TypeAnnInfo(:foo, mtsig, :( Tuple{Int} ))
     )
+
+    @test collectTypeAnnotations(fdefArgs) == list(
+        TypeAnnInfo(NOTAFUNSIG, tyassorann, :( Vector{<:Int} )),
+        TypeAnnInfo(:fdefArgs, mtsig, :( Tuple{Any, Vector{Bool}, Int} ))
+    )
 end
 
-# FIXME: collectFunDefTypeAnnotations currently extracts only method signature
+# FIXME: collectFunDefTypeSignature currently extracts only method signature
 @testset "types-analysis :: collect all ty-anns in file " begin
     @test parseAndCollectTypeAnnotations(testFilePath("empty.jl")) == nil()
 
     @test parseAndCollectTypeAnnotations(
             testFilePath("Multisets-cut.jl")
         ) == list(
             TypeAnnInfo(:push!, mtsig, :( Tuple{Multiset{T}, Any, Int} where T )),
+            TypeAnnInfo(:getindex, retty, :Int), 
             TypeAnnInfo(:getindex, mtsig, :( Tuple{Multiset{T}, Any} where T )),
             TypeAnnInfo(:clean!, mtsig, :( Tuple{Multiset} )),
             TypeAnnInfo(:Multiset, mtsig, :( Tuple{Base.AbstractSet{T}} where T )),
@@ -192,7 +213,8 @@ end
             TypeAnnInfo(:(Base.copy), mtsig, :( Tuple{Multiset{T}} where T )),
             TypeAnnInfo(:Multiset, mtsig, :( Tuple{Vararg{Any}} )),
             TypeAnnInfo(:Multiset, mtsig, :( Tuple{} )),
-            TypeAnnInfo(:Multiset, mtsig, :( Tuple{} where T ))
+            TypeAnnInfo(:Multiset, mtsig, :( Tuple{} where T )),
+            TypeAnnInfo(NOTAFUNSIG, tyassorann, :( Dict{T,Int} ))
         )
 end