entc: add option to override codegen header

Summary:
Set the standard header ("Code generated by entc, DO NOT EDIT.") as default, and override it using option in graph.

No changes to graph except the `generate.go` file.

Reviewed By: idoshveki

Differential Revision: D16642348

fbshipit-source-id: d9fd1d2046e2fd96acbb100ef061fda75d99ce52

entc/gen/func.go

@@ -38,6 +38,7 @@ var (
 		"join":        join,
 		"lower":       strings.ToLower,
 		"upper":       strings.ToUpper,
+		"hasField":    hasField,
 		"hasSuffix":   strings.HasSuffix,
 		"xtemplate":   xtemplate,
 		"hasTemplate": hasTemplate,
@@ -143,27 +144,40 @@ func receiver(s string) (r string) {
 	return strings.ToLower(s)
 }
 
-// scope wraps the Type object with extended context.
-type scope struct {
+// typeScope wraps the Type object with extended scope.
+type typeScope struct {
 	*Type
 	Scope map[interface{}]interface{}
 }
 
+// graphScope wraps the Graph object with extended scope.
+type graphScope struct {
+	*Graph
+	Scope map[interface{}]interface{}
+}
+
 // extend extends the parent block with a KV pairs.
 //
 //	{{ with $scope := extend $ "key" "value" }}
 //		{{ template "setters" $scope }}
 //	{{ end}}
 //
-func extend(t *Type, kv ...interface{}) *scope {
-	s := &scope{Type: t, Scope: make(map[interface{}]interface{})}
+func extend(v interface{}, kv ...interface{}) (interface{}, error) {
+	scope := make(map[interface{}]interface{})
 	if len(kv)%2 != 0 {
-		panic("invalid number of parameters")
+		return nil, fmt.Errorf("invalid number of parameters: %d", len(kv))
 	}
 	for i := 0; i < len(kv); i += 2 {
-		s.Scope[kv[i]] = kv[i+1]
+		scope[kv[i]] = kv[i+1]
+	}
+	switch v := v.(type) {
+	case *Type:
+		return &typeScope{Type: v, Scope: scope}, nil
+	case *Graph:
+		return &graphScope{Graph: v, Scope: scope}, nil
+	default:
+		return nil, fmt.Errorf("invalid type for extend: %T", v)
 	}
-	return s
 }
 
 // add calculates summarize list of variables.
@@ -250,3 +264,9 @@ func hasTemplate(name string) bool {
 	}
 	return false
 }
+
+// hasField determines if a struct has a field with the given name.
+func hasField(v interface{}, name string) bool {
+	vr := reflect.Indirect(reflect.ValueOf(v))
+	return vr.FieldByName(name).IsValid()
+}

entc/gen/template/base.tmpl

@@ -1,7 +1,7 @@
 {{ define "base" }}
 
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/builder/create.tmpl

@@ -1,6 +1,7 @@
 {{ define "create" }}
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/builder/delete.tmpl

@@ -1,6 +1,7 @@
 {{ define "delete" }}
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/builder/query.tmpl

@@ -1,6 +1,7 @@
 {{ define "query" }}
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/builder/update.tmpl

@@ -1,6 +1,7 @@
 {{ define "update" }}
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/client.tmpl

@@ -1,7 +1,7 @@
 {{ define "client" }}
 
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 import (
 	"log"

entc/gen/template/config.tmpl

@@ -1,7 +1,7 @@
 {{ define "config" }}
 
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/context.tmpl

@@ -1,7 +1,7 @@
 {{ define "context" }}
 
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 import (
 	"context"

entc/gen/template/ent.tmpl

@@ -1,7 +1,7 @@
 {{ define "model" }}
 
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 {{ template "import" $ }}
 

entc/gen/template/example.tmpl

@@ -1,6 +1,7 @@
 {{ define "example" }}
+
 {{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 import (
 	"log"

entc/gen/template/header.tmpl

@@ -1,5 +1,9 @@
 {{ define "header" }}
-// Code generated (@generated) by entc, DO NOT EDIT.
+// {{ with $.Header }}{{ . }}{{ else }}Code generated by entc, DO NOT EDIT.{{ end }}
 
-package {{ $ }}
+{{ $pkg := base $.Config.Package }}
+{{ if hasField $ "Scope" }}
+	{{ $pkg = $.Scope.Package }}
+{{ end }}
+package {{ $pkg }}
 {{ end }}

entc/gen/template/meta.tmpl

@@ -1,6 +1,8 @@
 {{ define "meta" }}
 
-{{ template "header" $.Package }}
+{{- with extend $ "Package" $.Package -}}
+	{{ template "header" . }}
+{{ end }}
 
 {{ template "import" $ }}
 

entc/gen/template/migrate/migrate.tmpl

@@ -1,6 +1,8 @@
 {{ define "migrate" }}
-{{ template "header" "migrate" }}
-{{ $pkg := base $.Config.Package }}
+
+{{- with extend $ "Package" "migrate" -}}
+	{{ template "header" . }}
+{{ end }}
 
 import (
 	"context"

entc/gen/template/migrate/schema.tmpl

@@ -1,5 +1,8 @@
 {{ define "schema" }}
-{{ template "header" "migrate" }}
+
+{{- with extend $ "Package" "migrate" -}}
+	{{ template "header" . }}
+{{ end }}
 
 import (
 	"fbc/ent/field"

entc/gen/template/predicate.tmpl

@@ -1,6 +1,8 @@
 {{ define "predicate" }}
 
-{{ template "header" "predicate" }}
+{{- with extend $ "Package" "predicate" -}}
+	{{ template "header" . }}
+{{ end }}
 
 import (
 	"fmt"

entc/gen/template/tx.tmpl

@@ -1,7 +1,6 @@
 {{ define "tx" }}
 
-{{ $pkg := base $.Config.Package }}
-{{ template "header" $pkg }}
+{{ template "header" $ }}
 
 import (
 	"context"

entc/gen/template/where.tmpl

@@ -1,6 +1,8 @@
 {{ define "where" }}
 
-{{ template "header" $.Package }}
+{{- with extend $ "Package" $.Package -}}
+	{{ template "header" . }}
+{{ end }}
 
 {{ template "import" $ }}
 

entc/integration/generate.go

@@ -1,5 +1,5 @@
 package integration
 
-//go:generate go run ../cmd/entc/entc.go generate --storage=sql,gremlin --idtype=string ./ent/schema
-//go:generate go run ../cmd/entc/entc.go generate ./migrate/entv1/schema
-//go:generate go run ../cmd/entc/entc.go generate ./migrate/entv2/schema
+//go:generate go run ../cmd/entc/entc.go generate --storage=sql,gremlin --idtype string --header "Code generated (@generated) by entc, DO NOT EDIT." ./ent/schema
+//go:generate go run ../cmd/entc/entc.go generate --header "Code generated (@generated) by entc, DO NOT EDIT." ./migrate/entv1/schema
+//go:generate go run ../cmd/entc/entc.go generate --header "Code generated (@generated) by entc, DO NOT EDIT." ./migrate/entv2/schema