import List + Dict + group typings

GetStream · Sep 18, 2023 · 2ed1bb8 · 2ed1bb8
1 parent 83887c8
commit 2ed1bb8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/openapi-gen/builtin_func.go b/openapi-gen/builtin_func.go
@@ -38,6 +38,21 @@ func PrepareBuiltinFunctions(config *Config) template.FuncMap {
 		"toCamel":   strcase.ToCamel,
 		"lower":     strings.ToLower,
 		"newSet":    newSet,
+		"append": func(value []string, elems ...string) []string {
+			return append(value, elems...)
+		},
+		"list": func(value ...string) []string {
+			return value
+		},
+		"contains": func(value []string, elem string) bool {
+			for _, el := range value {
+				if el == elem {
+					return true
+				}
+			}
+			return false
+		},
+		"join": strings.Join,
 		"has": func(sl []string, str string) bool {
 			for _, s := range sl {
 				if s == str {
@@ -46,7 +61,7 @@ func PrepareBuiltinFunctions(config *Config) template.FuncMap {
 			}
 			return false
 		},
-		"toUpper":        strings.ToUpper,
+		"toUpper":    strings.ToUpper,
 		"toConstant": toConstant,
 		"successfulResponse": func(responses openapi3.Responses) *openapi3.SchemaRef {
 			for code, response := range responses {

diff --git a/openapi-gen/templates/python/type.tmpl b/openapi-gen/templates/python/type.tmpl
@@ -1,10 +1,21 @@
 {{- $required := .Schema.Required}}
 {{- $hasDateTime := false }}
+{{- $typingImports := list}}
 {{- $imports := newSet}}
 {{- range $index, $value := sortedProperties .Schema}}
     {{- if and (eq .Value.Type "string") (eq .Value.Format "date-time")}}
         {{- $hasDateTime = true }}
     {{- end}}
+    {{- if eq .Value.Type "array"}}
+        {{- if not (contains $typingImports "List")}}
+            {{- $typingImports = append $typingImports "List"}}
+        {{- end}}
+    {{- end}}
+    {{- if and (eq .Value.Type "object") (.Value.AdditionalProperties) (.Value.AdditionalProperties.Schema)}}
+        {{- if not (contains $typingImports "Dict")}}
+            {{- $typingImports = append $typingImports "Dict"}}
+        {{- end}}
+    {{- end}}
 {{- end}}
 
 {{- if (not .Schema.Enum)}}
@@ -15,7 +26,12 @@ from enum import Enum
 {{- end}}
 
 {{- if .HasNonRequired}}
-from typing import Optional
+    {{- if not (contains $typingImports "Optional") -}}
+        {{- $typingImports = append $typingImports "Optional" -}}
+    {{- end}}
+{{- end}}
+{{- if gt (len $typingImports) 0}}
+from typing import {{join $typingImports ", "}}
 {{- end}}
 
 {{- if $hasDateTime}}
@@ -44,8 +60,8 @@ from {{with index additionalParameters "modelImportPrefix"}}{{.}}{{else}}models.
 {{- if (eq .Type "integer")}}int{{end}}
 {{- if (eq .Type "number")}}float{{end}}
 {{- if (eq .Type "boolean")}}bool{{end}}
-{{- if (eq .Type "array")}}list[{{template "generateSchemaRef" .Items -}}]{{end}}
-{{- if (eq .Type "object")}}{{with .AdditionalProperties}}dict[str, {{template "generateSchemaRef" .Schema -}}]{{end}}{{end}}
+{{- if (eq .Type "array")}}List[{{template "generateSchemaRef" .Items -}}]{{end}}
+{{- if (eq .Type "object")}}{{with .AdditionalProperties}}Dict[str, {{template "generateSchemaRef" .Schema -}}]{{end}}{{end}}
 {{- if (eq .Type "null")}}None{{end}}
 {{- if (eq .Type "")}}object{{end}}
 {{- end -}}