diff --git a/pkg/codegen/python/gen.go b/pkg/codegen/python/gen.go index 167c1840a732..3ca743c4b76c 100644 --- a/pkg/codegen/python/gen.go +++ b/pkg/codegen/python/gen.go @@ -357,27 +357,52 @@ func genStandardHeader(w io.Writer, tool string) { fmt.Fprintf(w, "# *** Do not edit by hand unless you're certain you know what you are doing! ***\n\n") } +func typingImports() []string { + return []string{ + "Any", + "Mapping", + "Optional", + "Sequence", + "Union", + "overload", + } +} + +func (mod *modContext) generateCommonImports(w io.Writer, imports imports, typingImports []string) { + rel, err := filepath.Rel(mod.mod, "") + contract.Assert(err == nil) + relRoot := path.Dir(rel) + relImport := relPathToRelImport(relRoot) + + fmt.Fprintf(w, "import copy\n") + fmt.Fprintf(w, "import warnings\n") + fmt.Fprintf(w, "import pulumi\n") + fmt.Fprintf(w, "import pulumi.runtime\n") + fmt.Fprintf(w, "from typing import %s\n", strings.Join(typingImports, ", ")) + fmt.Fprintf(w, "from %s import _utilities\n", relImport) + for _, imp := range imports.strings() { + fmt.Fprintf(w, "%s\n", imp) + } + fmt.Fprintf(w, "\n") +} + func (mod *modContext) genHeader(w io.Writer, needsSDK bool, imports imports) { genStandardHeader(w, mod.tool) // If needed, emit the standard Pulumi SDK import statement. if needsSDK { - rel, err := filepath.Rel(mod.mod, "") - contract.Assert(err == nil) - relRoot := path.Dir(rel) - relImport := relPathToRelImport(relRoot) - - fmt.Fprintf(w, "import copy\n") - fmt.Fprintf(w, "import warnings\n") - fmt.Fprintf(w, "import pulumi\n") - fmt.Fprintf(w, "import pulumi.runtime\n") - fmt.Fprintf(w, "from typing import Any, Mapping, Optional, Sequence, Union, overload\n") - fmt.Fprintf(w, "from %s import _utilities\n", relImport) - for _, imp := range imports.strings() { - fmt.Fprintf(w, "%s\n", imp) - } - fmt.Fprintf(w, "\n") + typings := typingImports() + mod.generateCommonImports(w, imports, typings) + } +} + +func (mod *modContext) genFunctionHeader(w io.Writer, function *schema.Function, imports imports) { + genStandardHeader(w, mod.tool) + typings := typingImports() + if function.Outputs == nil || len(function.Outputs.Properties) == 0 { + typings = append(typings, "Awaitable") } + mod.generateCommonImports(w, imports, typings) } func relPathToRelImport(relPath string) string { @@ -1671,7 +1696,7 @@ func (mod *modContext) genFunction(fun *schema.Function) (string, error) { mod.collectImports(fun.Outputs.Properties, imports, false) } - mod.genHeader(w, true /*needsSDK*/, imports) + mod.genFunctionHeader(w, fun, imports) var baseName, awaitableName string if fun.Outputs != nil && len(fun.Outputs.Properties) > 0 { @@ -1702,6 +1727,8 @@ func (mod *modContext) genFunction(fun *schema.Function) (string, error) { if fun.Outputs != nil && len(fun.Outputs.Properties) > 0 { retTypeName, rets = mod.genAwaitableType(w, fun.Outputs), fun.Outputs.Properties fmt.Fprintf(w, "\n\n") + } else { + retTypeName = "Awaitable[None]" } var args []*schema.Property diff --git a/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_const_input.py b/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_const_input.py index e65015b2dd67..730d3c2fb85b 100644 --- a/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_const_input.py +++ b/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_const_input.py @@ -6,7 +6,7 @@ import warnings import pulumi import pulumi.runtime -from typing import Any, Mapping, Optional, Sequence, Union, overload +from typing import Any, Mapping, Optional, Sequence, Union, overload, Awaitable from . import _utilities __all__ = [ @@ -14,7 +14,7 @@ ] def func_with_const_input(plain_input: Optional[str] = None, - opts: Optional[pulumi.InvokeOptions] = None): + opts: Optional[pulumi.InvokeOptions] = None) -> Awaitable[None]: """ Codegen demo with const inputs """ diff --git a/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_empty_outputs.py b/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_empty_outputs.py index f1a3a6af86bb..9d0d6c81f9c4 100644 --- a/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_empty_outputs.py +++ b/pkg/codegen/testing/test/testdata/output-funcs/python/pulumi_mypkg/func_with_empty_outputs.py @@ -6,7 +6,7 @@ import warnings import pulumi import pulumi.runtime -from typing import Any, Mapping, Optional, Sequence, Union, overload +from typing import Any, Mapping, Optional, Sequence, Union, overload, Awaitable from . import _utilities __all__ = [ @@ -14,7 +14,7 @@ ] def func_with_empty_outputs(name: Optional[str] = None, - opts: Optional[pulumi.InvokeOptions] = None): + opts: Optional[pulumi.InvokeOptions] = None) -> Awaitable[None]: """ n/a diff --git a/pkg/codegen/testing/test/testdata/regress-node-8110/python/pulumi_my8110/example_func.py b/pkg/codegen/testing/test/testdata/regress-node-8110/python/pulumi_my8110/example_func.py index eb63bf80019c..f7e64ab757f1 100644 --- a/pkg/codegen/testing/test/testdata/regress-node-8110/python/pulumi_my8110/example_func.py +++ b/pkg/codegen/testing/test/testdata/regress-node-8110/python/pulumi_my8110/example_func.py @@ -6,7 +6,7 @@ import warnings import pulumi import pulumi.runtime -from typing import Any, Mapping, Optional, Sequence, Union, overload +from typing import Any, Mapping, Optional, Sequence, Union, overload, Awaitable from . import _utilities from ._enums import * @@ -15,7 +15,7 @@ ] def example_func(enums: Optional[Sequence[Union[str, 'MyEnum']]] = None, - opts: Optional[pulumi.InvokeOptions] = None): + opts: Optional[pulumi.InvokeOptions] = None) -> Awaitable[None]: """ Use this data source to access information about an existing resource. """ diff --git a/pkg/codegen/testing/test/testdata/simple-plain-schema/python/pulumi_example/do_foo.py b/pkg/codegen/testing/test/testdata/simple-plain-schema/python/pulumi_example/do_foo.py index 3635cb794612..17cf22139d81 100644 --- a/pkg/codegen/testing/test/testdata/simple-plain-schema/python/pulumi_example/do_foo.py +++ b/pkg/codegen/testing/test/testdata/simple-plain-schema/python/pulumi_example/do_foo.py @@ -6,7 +6,7 @@ import warnings import pulumi import pulumi.runtime -from typing import Any, Mapping, Optional, Sequence, Union, overload +from typing import Any, Mapping, Optional, Sequence, Union, overload, Awaitable from . import _utilities from ._inputs import * @@ -15,7 +15,7 @@ ] def do_foo(foo: Optional[pulumi.InputType['Foo']] = None, - opts: Optional[pulumi.InvokeOptions] = None): + opts: Optional[pulumi.InvokeOptions] = None) -> Awaitable[None]: """ Use this data source to access information about an existing resource. """