ent/entc/gen/template/builder/setter.tmpl

{{/*
Copyright 2019-present Facebook Inc. All rights reserved.
This source code is licensed under the Apache 2.0 license found
in the LICENSE file in the root directory of this source tree.
*/}}

{{ define "setter" }}
{{ $builder := pascal $.Scope.Builder }}
{{ $receiver := receiver $builder }}
{{ $fields := $.Fields }}
{{ $updater := false }}
{{- if or (hasSuffix $builder "Update") (hasSuffix $builder "UpdateOne") }}
	{{ $updater = true }}
	{{ $fields = $.MutableFields }}
{{- end }}

{{ range $_, $f := $fields }}
	{{ $p := receiver $f.Type.String }}
	{{ $func := print "Set" $f.StructField }}
	// {{ $func }} sets the {{ $f.Name }} field.
	func ({{ $receiver }} *{{ $builder }}) {{ $func }}({{ $p }} {{ $f.Type }}) *{{ $builder }} {
		{{ $receiver }}.{{ $f.BuilderField }} = &{{ $p }}
		{{- /* setting numeric type override previous calls to Add. */}}
		{{- if and $f.Type.Numeric $updater }}
			{{ $receiver }}.add{{ $f.BuilderField }} = nil
		{{- end }}
		return {{ $receiver }}
	}

	{{/* avoid generting nillable setters for nillable types. */}}
	{{ if and (not $f.Type.Nillable) (or $f.Optional $f.Default) (not (and $updater $f.UpdateDefault)) }}
		{{ $nillableFunc := print "SetNillable" $f.StructField }}
		// {{ $nillableFunc }} sets the {{ $f.Name }} field if the given value is not nil.
		func ({{ $receiver }} *{{ $builder }}) {{ $nillableFunc }}({{ $p }} *{{ $f.Type }}) *{{ $builder }} {
			if {{ $p }} != nil {
				{{ $receiver }}.{{ $func }}(*{{ $p }})
			}
			return {{ $receiver }}
		}
	{{ end }}

	{{ if and $f.Type.Numeric $updater }}
		{{ $func := print "Add" $f.StructField }}
		// {{ $func }} adds {{ $p }} to {{ $f.Name }}.
		func ({{ $receiver }} *{{ $builder }}) {{ $func }}({{ $p }} {{ $f.Type }}) *{{ $builder }} {
			if {{ $receiver }}.add{{ $f.BuilderField }} == nil {
				{{ $receiver }}.add{{ $f.BuilderField }} = &{{ $p }}
			} else {
				*{{ $receiver }}.add{{ $f.BuilderField }} += {{ $p }}
			}
			return {{ $receiver }}
		}
	{{ end }}

	{{ if and $f.Optional $updater }}
		{{ $func := print "Clear" $f.StructField }}
		// {{ $func }} clears the value of {{ $f.Name }}.
		func ({{ $receiver }} *{{ $builder }}) {{ $func }}() *{{ $builder }} {
			{{ $receiver }}.{{ $f.BuilderField }} = nil
			{{ $receiver }}.clear{{ $f.BuilderField }} = true
			return {{ $receiver }}
		}
	{{ end }}
{{ end }}

{{ range $_, $e := $.Edges }}
	{{ $op := "add" }}{{ if $e.Unique }}{{ $op = "set" }}{{ end }}
	{{ $idsFunc := print (pascal $op) (singular $e.Name | pascal) "IDs" }}{{ if $e.Unique }}{{ $idsFunc = print (pascal $op) (pascal $e.Name) "ID" }}{{ end }}
	// {{ $idsFunc }} {{ $op }}s the {{ $e.Name }} edge to {{ $e.Type.Name }} by id{{ if not $e.Unique }}s{{ end }}.
	func ({{ $receiver }} *{{ $builder }}) {{ $idsFunc }}({{ if $e.Unique }}id{{ else }}ids ...{{ end }} {{ $.ID.Type }}) *{{ $builder }} {
		if {{ $receiver }}.{{ $e.BuilderField }} == nil {
			{{ $receiver }}.{{ $e.BuilderField }} = make(map[{{ $.ID.Type }}]struct{})
		}
		{{ if $e.Unique -}}
			{{ $receiver }}.{{ $e.BuilderField }}[id] = struct{}{}
		{{- else -}}
			for i := range ids {
				{{ $receiver }}.{{ $e.BuilderField }}[ids[i]] = struct{}{}
			}
		{{- end }}
		return {{ $receiver }}
	}
	{{ if and $e.Unique $e.Optional }}
		{{ $nillableIDsFunc := print "SetNillable" $e.StructField "ID" }}
		// {{ $nillableIDsFunc }} sets the {{ $e.Name }} edge to {{ $e.Type.Name }} by id if the given value is not nil.
		func ({{ $receiver }} *{{ $builder }}) {{ $nillableIDsFunc }}(id *{{ $.ID.Type }}) *{{ $builder }} {
			if id != nil {
				{{ $receiver}} = {{ $receiver }}.{{ $idsFunc }}(*id)
			}
			return {{ $receiver }}
		}
	{{ end }}
	{{ $p := lower (printf "%.1s" $e.Type.Name) }}
	{{ if eq $p $receiver }} {{ $p = "v" }} {{ end }}
	{{ $func := print (pascal $op) $e.StructField }}
	// {{ $func }} {{ $op }}s the {{ $e.Name }} edge{{if not $e.Unique}}s{{ end }} to {{ $e.Type.Name }}.
	func ({{ $receiver }} *{{ $builder }}) {{ $func }}({{ $p }} {{ if not $e.Unique }}...{{ end }}*{{ $e.Type.Name}}) *{{ $builder }} {
		{{ if $e.Unique -}}
			return {{ $receiver }}.{{ $idsFunc }}({{ $p }}.ID)
		{{- else -}}
			ids := make([]{{ $.ID.Type }}, len({{ $p }}))
			{{ $i := "i" }}{{ if eq $i $p }}{{ $i = "j" }}{{ end -}}
			for {{ $i }} := range {{ $p }} {
				ids[{{ $i }}] = {{ $p }}[{{ $i }}].ID
			}
			return {{ $receiver }}.{{ $idsFunc }}(ids...)
		{{- end }}
	}
{{ end }}
{{ end }}