diff --git a/internal/jennies/python/builder.go b/internal/jennies/python/builder.go index 909ea8255..d6774c388 100644 --- a/internal/jennies/python/builder.go +++ b/internal/jennies/python/builder.go @@ -87,8 +87,9 @@ func (jenny *Builder) generateBuilder(context languages.Context, builder ast.Bui return jenny.tmpl. Funcs(common.TypeResolvingTemplateHelpers(context)). Funcs(map[string]any{ - "formatType": jenny.typeFormatter.formatType, - "formatRawType": jenny.rawTypeFormatter.formatType, + "isDisjunctionOfBuilders": context.IsDisjunctionOfBuilders, + "formatType": jenny.typeFormatter.formatType, + "formatRawType": jenny.rawTypeFormatter.formatType, "formatRawTypeNotNullable": func(def ast.Type) string { typeDef := def.DeepCopy() typeDef.Nullable = false diff --git a/internal/jennies/python/templates/builders/assignment.tmpl b/internal/jennies/python/templates/builders/assignment.tmpl index 3a171f9f6..f698448c4 100644 --- a/internal/jennies/python/templates/builders/assignment.tmpl +++ b/internal/jennies/python/templates/builders/assignment.tmpl @@ -21,7 +21,12 @@ {{- with .Value.Argument }} {{- if or (resolvesToBuilder .Type) (resolvesToComposableSlot .Type) }} {{- $builtResultSuffix := ternary "_resources" "_resource" (or .Type.IsArray .Type.IsMap) }} -{{ .Name|formatIdentifier }}{{ $builtResultSuffix }} = {{ template "unfold_builders" (dict "InputType" .Type "InputVar" (formatIdentifier .Name) "Depth" 1) }} +{{- $name := print (.Name|formatIdentifier) $builtResultSuffix }} +{{ $name }} = {{ template "unfold_builders" (dict "InputType" .Type "InputVar" (formatIdentifier .Name) "Depth" 1) }} +{{- if and .Type.IsDisjunction (not (isDisjunctionOfBuilders .Type)) }} +{{ $branchesCount := sub1 (len .Type.Disjunction.Branches) -}} +assert {{ range $i, $branch := .Type.Disjunction.Branches }}isinstance({{ $name }}, {{ $branch|formatRawType}}){{ if ne $i $branchesCount }} or {{ end }}{{ end -}} +{{- end -}} {{- end }} {{- end }} {{- with .Value.Envelope }} @@ -36,6 +41,8 @@ [{{ template "unfold_builders" (dict "InputType" .InputType.Array.ValueType "InputVar" (print "r" .Depth ) "Depth" (add1 .Depth)) }} for r{{ .Depth }} in {{ .InputVar }}] {{- else if .InputType.IsMap -}} { key{{.Depth}}: {{ template "unfold_builders" (dict "InputType" .InputType.Map.ValueType "InputVar" (print "val" .Depth ) "Depth" (add1 .Depth)) }} for (key{{.Depth}}, val{{ .Depth }}) in {{ .InputVar }}.items() } +{{- else if and .InputType.IsDisjunction (not (isDisjunctionOfBuilders .InputType)) -}} +{{ .InputVar }}.build() if hasattr({{ .InputVar }}, 'build') and callable({{ .InputVar }}.build) else {{ .InputVar }} {{- else -}} {{ .InputVar }}.build() {{- end -}} diff --git a/internal/jennies/python/tmpl.go b/internal/jennies/python/tmpl.go index 16eb40805..8683cf221 100644 --- a/internal/jennies/python/tmpl.go +++ b/internal/jennies/python/tmpl.go @@ -24,6 +24,9 @@ func initTemplates(apiRefCollector *common.APIReferenceCollector, extraTemplates template.Funcs(formattingTemplateFuncs()), // placeholder functions, will be overridden by jennies template.Funcs(template.FuncMap{ + "isDisjunctionOfBuilders": func(_ ast.Type) string { + panic("isDisjunctionOfBuilders() needs to be overridden by a jenny") + }, "formatType": func(_ ast.Type) string { panic("formatType() needs to be overridden by a jenny") }, diff --git a/testdata/jennies/builders/builder_delegation_in_disjunction/PythonBuilder/builders/builder_delegation_in_disjunction.py b/testdata/jennies/builders/builder_delegation_in_disjunction/PythonBuilder/builders/builder_delegation_in_disjunction.py index fbb64a265..9d02dfc23 100644 --- a/testdata/jennies/builders/builder_delegation_in_disjunction/PythonBuilder/builders/builder_delegation_in_disjunction.py +++ b/testdata/jennies/builders/builder_delegation_in_disjunction/PythonBuilder/builders/builder_delegation_in_disjunction.py @@ -61,7 +61,8 @@ def single_link_or_string(self, single_link_or_string: typing.Union[cogbuilder.B will be expanded to cog.Builder | string """ - single_link_or_string_resource = single_link_or_string.build() + single_link_or_string_resource = single_link_or_string.build() if hasattr(single_link_or_string, 'build') and callable(single_link_or_string.build) else single_link_or_string + assert isinstance(single_link_or_string_resource, builder_delegation_in_disjunction.DashboardLink) or isinstance(single_link_or_string_resource, str) self._internal.single_link_or_string = single_link_or_string_resource return self @@ -71,7 +72,7 @@ def links_or_strings(self, links_or_strings: list[typing.Union[cogbuilder.Builde will be expanded to [](cog.Builder | string) """ - links_or_strings_resources = [r1.build() for r1 in links_or_strings] + links_or_strings_resources = [r1.build() if hasattr(r1, 'build') and callable(r1.build) else r1 for r1 in links_or_strings] self._internal.links_or_strings = links_or_strings_resources return self