Implement definition(method, String) and add internal data

Author: timholy

src/CodeTracking.jl

@@ -1,15 +1,80 @@
 module CodeTracking
 
-export whereis
+using Core: LineInfoNode
 
-# This is just a stub implementation for now
+export whereis, definition, pkgfiles
+
+include("data.jl")
+include("utils.jl")
+
+"""
+    filepath, line = whereis(method::Method)
+
+Return the file and line of the definition of `method`. `line`
+is the first line of the method's body.
+"""
 function whereis(method::Method)
-    file, line = String(method.file), method.line
+    lin = get(method_locations, method.sig, nothing)
+    if lin === nothing
+        file, line = String(method.file), method.line
+    else
+        file, line = fileline(lin)
+    end
     if !isabspath(file)
-        # This is a Base method
+        # This is a Base or Core method
         file = Base.find_source_file(file)
     end
     return normpath(file), line
 end
 
+"""
+    src = definition(method::Method, String)
+
+Return a string with the code that defines `method`.
+
+Note this may not be terribly useful for methods that are defined inside `@eval` statements;
+see [`definition(method::Method, Expr)`](@ref) instead.
+"""
+function definition(method::Method, ::Type{String})
+    file, line = whereis(method)
+    src = read(file, String)
+    eol = isequal('\n')
+    linestarts = Int[]
+    istart = 0
+    for i = 1:line-1
+        push!(linestarts, istart+1)
+        istart = findnext(eol, src, istart+1)
+    end
+    ex, iend = Meta.parse(src, istart)
+    if isfuncexpr(ex)
+        return src[istart+1:iend-1]
+    end
+    # The function declaration was presumably on a previous line
+    lineindex = lastindex(linestarts)
+    while !isfuncexpr(ex)
+        istart = linestarts[lineindex]
+        ex, iend = Meta.parse(src, istart)
+    end
+    return src[istart:iend-1]
+end
+
+"""
+    ex = definition(method::Method, Expr)
+    ex = definition(method::Method)
+
+Return an expression that defines `method`.
+"""
+definition(method::Method, ::Type{Expr}) = get(method_definitions, method.sig, nothing)
+
+definition(method::Method) = definition(method, Expr)
+
+"""
+    files = pkgfiles(mod::Module)
+
+Return a list of the files that were loaded to define `mod`.
+"""
+function pkgfiles(mod::Module)
+    error("not implemented")
+end
+
 end # module

src/data.jl

@@ -0,0 +1,5 @@
+# The variables here get populated by Revise.jl.
+
+const method_locations = IdDict{Type,LineInfoNode}()
+
+const method_definitions = IdDict{Type,Expr}()

src/utils.jl

@@ -0,0 +1,16 @@
+function isfuncexpr(ex)
+    ex.head == :function && return true
+    if ex.head == :(=)
+        a = ex.args[1]
+        if isa(a, Expr)
+            while a.head == :where
+                a = a.args[1]
+                isa(a, Expr) || return false
+            end
+            a.head == :call && return true
+        end
+    end
+    return false
+end
+
+fileline(lin::LineInfoNode) = String(lin.file), lin.line

test/runtests.jl

@@ -1,6 +1,22 @@
 using CodeTracking
 using Test
 
+include("script.jl")
+
 @testset "CodeTracking.jl" begin
-    # Write your own tests here.
+    m = first(methods(f1))
+    file, line = whereis(m)
+    @test file == normpath(joinpath(@__DIR__, "script.jl"))
+    src = definition(m, String)
+    @test src == """
+    function f1(x, y)
+        return x + y
+    end
+    """
+
+    m = first(methods(f2))
+    src = definition(m, String)
+    @test src == """
+    f2(x, y) = x + y
+    """
 end

test/script.jl

@@ -0,0 +1,5 @@
+function f1(x, y)
+    return x + y
+end
+
+f2(x, y) = x + y