[beam] branch master updated: [BEAM-6577] Examine func return types when generating shims (#7720)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

altay pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 3b139a7  [BEAM-6577] Examine func return types when generating shims (#7720)
3b139a7 is described below

commit 3b139a77387f9a097fb6f1bb0cabc7343ce810ac
Author: Tianyang Hu <ht...@gmail.com>
AuthorDate: Mon Feb 4 14:27:10 2019 -0800

    [BEAM-6577] Examine func return types when generating shims (#7720)
    
    * Examine func return types when generating shims
    
    starcgenx.Extractor should examine both function input and output when extracting types from a package. Otherwise, the generated code may miss certain package imports.
    
    * Refactor extractParameters
    
    Rename extractParameters to extractFromSignature, which calls extractFromTuple using both params and results.
---
 sdks/go/pkg/beam/util/starcgenx/starcgenx.go      | 24 +++++++++++++----------
 sdks/go/pkg/beam/util/starcgenx/starcgenx_test.go | 17 ++++++++++++++++
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/sdks/go/pkg/beam/util/starcgenx/starcgenx.go b/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
index 82cfb74..0d47622 100644
--- a/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
+++ b/sdks/go/pkg/beam/util/starcgenx/starcgenx.go
@@ -290,7 +290,7 @@ func (e *Extractor) fromObj(fset *token.FileSet, id *ast.Ident, obj types.Object
 		}
 
 		e.funcs[e.sigKey(sig)] = sig
-		e.extractParameters(sig)
+		e.extractFromSignature(sig)
 		e.Printf("\t%v\n", sig)
 	case *types.TypeName:
 		e.Printf("%s: %q is a type %v --- %T %v %v %v %v\n",
@@ -325,12 +325,16 @@ func (e *Extractor) extractType(ot *types.TypeName) {
 	e.types[name] = struct{}{}
 }
 
-// Examines the signature and extracts types of parameters for generating
-// necessary imports and emitter and iterator code.
-func (e *Extractor) extractParameters(sig *types.Signature) {
-	in := sig.Params() // *types.Tuple
-	for i := 0; i < in.Len(); i++ {
-		s := in.At(i) // *types.Var
+// Examines the signature and extracts types of parameters and results for
+// generating necessary imports and emitter and iterator code.
+func (e *Extractor) extractFromSignature(sig *types.Signature) {
+	e.extractFromTuple(sig.Params())
+	e.extractFromTuple(sig.Results())
+}
+
+func (e *Extractor) extractFromTuple(tuple *types.Tuple) {
+	for i := 0; i < tuple.Len(); i++ {
+		s := tuple.At(i) // *types.Var
 
 		// Pointer types need to be iteratively unwrapped until we're at the base type,
 		// so we can get the import if necessary.
@@ -340,7 +344,7 @@ func (e *Extractor) extractParameters(sig *types.Signature) {
 			t = p.Elem()
 			p, ok = t.(*types.Pointer)
 		}
-		// Here's were we ensure we register new imports.
+		// Here's where we ensure we register new imports.
 		if t, ok := t.(*types.Named); ok {
 			if pkg := t.Obj().Pkg(); pkg != nil {
 				e.imports[pkg.Path()] = struct{}{}
@@ -357,8 +361,8 @@ func (e *Extractor) extractParameters(sig *types.Signature) {
 			if ipt, ok := e.makeInput(a); ok {
 				e.iters[ipt.Name] = ipt
 			}
-			// Tail recurse on functional parameters.
-			e.extractParameters(a)
+			// Tail recurse on functional signature.
+			e.extractFromSignature(a)
 		}
 	}
 }
diff --git a/sdks/go/pkg/beam/util/starcgenx/starcgenx_test.go b/sdks/go/pkg/beam/util/starcgenx/starcgenx_test.go
index 4b437c1..db59882 100644
--- a/sdks/go/pkg/beam/util/starcgenx/starcgenx_test.go
+++ b/sdks/go/pkg/beam/util/starcgenx/starcgenx_test.go
@@ -36,6 +36,9 @@ func TestExtractor(t *testing.T) {
 		{name: "pardo1", files: []string{pardo}, pkg: "pardo",
 			expected: []string{"runtime.RegisterFunction(MyIdent)", "runtime.RegisterFunction(MyDropVal)", "runtime.RegisterFunction(MyOtherDoFn)", "runtime.RegisterType(reflect.TypeOf((*foo)(nil)).Elem())", "funcMakerStringГString", "funcMakerIntStringГInt", "funcMakerFooГStringFoo"},
 		},
+		{name: "imports1", files: []string{imports}, pkg: "imports",
+			expected: []string{"runtime.RegisterFunction(MyImportedTypesDoFn)", "runtime.RegisterType(reflect.TypeOf((*rand.Rand)(nil)).Elem())", "runtime.RegisterType(reflect.TypeOf((*time.Time)(nil)).Elem())", "funcMakerRand۰RandГTime۰Time", "\"math/rand\"", "\"time\""},
+		},
 		{name: "emits1", files: []string{emits}, pkg: "emits",
 			expected: []string{"runtime.RegisterFunction(anotherFn)", "runtime.RegisterFunction(emitFn)", "runtime.RegisterType(reflect.TypeOf((*reInt)(nil)).Elem())", "funcMakerEmitIntIntГ", "emitMakerIntInt", "funcMakerIntIntEmitIntIntГError"},
 		},
@@ -100,6 +103,20 @@ func MyOtherDoFn(v foo) (string,foo) {
 }
 `
 
+const imports = `
+package imports
+
+import (
+	"math/rand"
+	"time"
+)
+
+// Imported types
+func MyImportedTypesDoFn(r *rand.Rand) time.Time {
+	return time.Unix(r.Int63(), 0)
+}
+`
+
 const emits = `
 package emits