al3x
al3x

Reputation: 1846

How to parse the statements in a method or function with Go

I would like to parse go code, and more especially the content of a function.

So far it's pretty easy to get the function declaration with the parser.ParseFile function.

package main

import (
    "fmt"
    "go/parser"
    "go/token"
    "go/ast"
    "log"
)

var code = []byte(`
package main

import (
    "fmt"
)

func GetFoo() {
    test := foo()
    fmt.Println(test)
}

func foo() int {
    return 0
}
`)

func main() {
    fset := token.NewFileSet()
    f, err := parser.ParseFile(fset, "", code, parser.ParseComments)
    if err != nil {
        log.Fatal(err)
    }

    for _, decl := range f.Decls {
        switch t := decl.(type) {
        // That's a func decl !
        case *ast.FuncDecl:
            fmt.Printf("Function name: %v\n", t.Name)
            // Now I would like to access to the test var and get its type
            // Must be int, because foo() returns an int
        }
    }

}

Now I would like to access to the test var and get its type, but I am a little bit lost. I have seen the decl.Body.List to iterate over the statements but I can get that my test var is an int

Thank you for your precious help !

https://play.golang.org/p/Y8uwM-CDWy

Upvotes: 3

Views: 3455

Answers (1)

al3x
al3x

Reputation: 1846

Thanks to JimB and his hint about go/types, here's how I could get the signature of my method

package main

import (
    "fmt"
    "go/ast"
    "go/importer"
    "go/parser"
    "go/token"
    "go/types"
    "log"
)

var code = []byte(`
package main

import (
    "fmt"
)

func GetFoo() {
    test := foo()
    fmt.Println(test)
}

func foo() int {
    return 0
}
`)

func main() {
    fset := token.NewFileSet()
    f, err := parser.ParseFile(fset, "", code, parser.ParseComments)
    if err != nil {
        log.Fatal(err)
    }

    conf := types.Config{Importer: importer.Default()}
    pkg, err := conf.Check("cmd", fset, []*ast.File{f}, nil)
    scope := pkg.Scope()

    for _, decl := range f.Decls {
        switch t := decl.(type) {
        // That's a func decl !
        case *ast.FuncDecl:
            for _, s := range t.Body.List {
                switch as := s.(type) {
                case *ast.AssignStmt:
                    for _, l := range as.Rhs {
                        switch rh := l.(type) {
                        case *ast.CallExpr:
                            if i, ok := rh.Fun.(*ast.Ident); ok {
                                ft := scope.Lookup(i.Name)
                                if ft != nil {
                                    if ftype, ok := ft.(*types.Func); ok {
                                        f := ftype.Type()
                                        if sig, ok := f.(*types.Signature); ok {
                                            // get the returned elements
                                            r := sig.Results()
                                            for i := 0; i < r.Len(); i++ {
                                                v := r.At(i)
                                                varTpe := v.Type()
                                                fmt.Println(varTpe)
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

Upvotes: 2

Related Questions