[sdk/python] - Support enums (#5615)

Co-authored-by: Pat Gavlin <pat@pulumi.com>

CHANGELOG.md

@@ -3,6 +3,9 @@ CHANGELOG
 
 ## HEAD (Unreleased)
 
+- [sdk/python] Add deserialization support for enums.
+  [#5615](https://github.com/pulumi/pulumi/pull/5615)
+
 - Respect `PULUMI_PYTHON_CMD` in scripts.
   [#5782](https://github.com/pulumi/pulumi/pull/5782)
 

sdk/python/lib/pulumi/runtime/rpc.py

@@ -22,6 +22,7 @@ import functools
 import inspect
 from abc import ABC, abstractmethod
 from typing import List, Any, Callable, Dict, Mapping, Optional, Sequence, Set, TYPE_CHECKING, cast
+from enum import Enum
 
 from google.protobuf import struct_pb2
 from semver import VersionInfo as Version # type:ignore
@@ -55,6 +56,7 @@ _special_resource_sig = "5cf8f73096256a8f31e491e813e4eb8e"
 
 _INT_OR_FLOAT = six.integer_types + (float,)
 
+
 def isLegalProtobufValue(value: Any) -> bool:
     """
     Returns True if the given value is a legal Protobuf value as per the source at
@@ -62,6 +64,7 @@ def isLegalProtobufValue(value: Any) -> bool:
     """
     return value is None or isinstance(value, (bool, six.string_types, _INT_OR_FLOAT, dict, list))
 
+
 async def serialize_properties(inputs: 'Inputs',
                                property_deps: Dict[str, List['Resource']],
                                input_transformer: Optional[Callable[[str], str]] = None) -> struct_pb2.Struct:
@@ -233,6 +236,7 @@ async def serialize_property(value: 'Input[Any]',
 
     return value
 
+
 # pylint: disable=too-many-return-statements
 def deserialize_properties(props_struct: struct_pb2.Struct, keep_unknowns: Optional[bool] = None) -> Any:
     """
@@ -290,6 +294,7 @@ def deserialize_properties(props_struct: struct_pb2.Struct, keep_unknowns: Optio
 
     return output
 
+
 def deserialize_resource(ref_struct: struct_pb2.Struct, keep_unknowns: Optional[bool] = None) -> 'Resource':
     urn = ref_struct["urn"]
     version = ref_struct["packageVersion"] if "packageVersion" in ref_struct else ""
@@ -321,12 +326,14 @@ def deserialize_resource(ref_struct: struct_pb2.Struct, keep_unknowns: Optional[
 
     return urn
 
+
 def is_rpc_secret(value: Any) -> bool:
     """
     Returns if a given python value is actually a wrapped secret.
     """
     return isinstance(value, dict) and _special_sig_key in value and value[_special_sig_key] == _special_secret_sig
 
+
 def wrap_rpc_secret(value: Any) -> Any:
     """
     Given a value, wrap it as a secret value if it isn't already a secret, otherwise return the value unmodified.
@@ -339,6 +346,7 @@ def wrap_rpc_secret(value: Any) -> Any:
         "value": value,
     }
 
+
 def unwrap_rpc_secret(value: Any) -> Any:
     """
     Given a value, if it is a wrapped secret value, return the underlying, otherwise return the value unmodified.
@@ -348,6 +356,7 @@ def unwrap_rpc_secret(value: Any) -> Any:
 
     return value
 
+
 def deserialize_property(value: Any, keep_unknowns: Optional[bool] = None) -> Any:
     """
     Deserializes a single protobuf value (either `Struct` or `ListValue`) into idiomatic
@@ -468,6 +477,8 @@ def translate_output_properties(output: Any,
 
     If output is a `float` and `typ` is `int`, the value is cast to `int`.
 
+    If output is in [`str`, `int`, `float`] and `typ` is an enum type, instantiate the enum type.
+
     Otherwise, if output is a primitive (i.e. not a dict or list), the value is returned without modification.
 
     :param Optional[type] typ: The output's target type.
@@ -497,7 +508,7 @@ def translate_output_properties(output: Any,
             # If typ is an output type, get its types, so we can pass
             # the type along for each property.
             types = _types.output_type_types(typ)
-            get_type = lambda k: types.get(k) # pylint: disable=unnecessary-lambda
+            get_type = lambda k: types.get(k)  # pylint: disable=unnecessary-lambda
         elif typ:
             # If typ is a dict, get the type for its values, to pass
             # along for each key.
@@ -540,6 +551,9 @@ def translate_output_properties(output: Any,
                 raise AssertionError(f"Unexpected type. Expected 'list' got '{typ}'")
         return [translate_output_properties(v, output_transformer, element_type) for v in output]
 
+    if typ and isinstance(output, (int, float, str)) and inspect.isclass(typ) and issubclass(typ, Enum):
+        return typ(output)
+
     if isinstance(output, float) and typ is int:
         return int(output)
 
@@ -572,7 +586,7 @@ def resolve_outputs(res: 'Resource',
     # outputs.  If the same property exists in the inputs and outputs states, the output wins.
     all_properties = {}
     # Get the resource's output types, so we can convert dicts from the engine into actual
-    # instantiated output types as needed.
+    # instantiated output types or primitive types into enums as needed.
     types = _types.resource_types(type(res))
     for key, value in deserialize_properties(outputs).items():
         # Outputs coming from the provider are NOT translated. Do so here.
@@ -592,6 +606,7 @@ def resolve_outputs(res: 'Resource',
 
     resolve_properties(resolvers, all_properties, deps)
 
+
 def resolve_properties(resolvers: Dict[str, Resolver], all_properties: Dict[str, Any], deps: Mapping[str, Set['Resource']]):
     for key, value in all_properties.items():
         # Skip "id" and "urn", since we handle those specially.
@@ -648,7 +663,7 @@ def resolve_outputs_due_to_exception(resolvers: Dict[str, Resolver], exn: Except
     failed to resolve.
 
     :param resolvers: Resolvers associated with a resource's outputs.
-    :param exn: The exception that occured when trying (and failing) to create this resource.
+    :param exn: The exception that occurred when trying (and failing) to create this resource.
     """
     for key, resolve in resolvers.items():
         log.debug(f"sending exception to resolver for {key}")

sdk/python/lib/test/test_next_serialize.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import asyncio
 import unittest
+from enum import Enum
 from typing import Any, Dict, List, Mapping, Optional, Sequence
 
 from google.protobuf import struct_pb2
@@ -222,7 +223,7 @@ class NextSerializationTests(unittest.TestCase):
 
         other = TestCustomResource("some-other-resource")
         other_fut = asyncio.Future()
-        other_fut.set_result(UNKNOWN) # <- not known
+        other_fut.set_result(UNKNOWN)  # <- not known
         other_known_fut = asyncio.Future()
         other_known_fut.set_result(False)
         other_out = Output({other}, other_fut, other_known_fut)
@@ -236,7 +237,6 @@ class NextSerializationTests(unittest.TestCase):
         # create it were unknown.
         self.assertEqual(rpc.UNKNOWN, prop)
 
-
     @async_test
     async def test_unknown_output(self):
         res = TestCustomResource("some-dependency")
@@ -835,7 +835,7 @@ class NextSerializationTests(unittest.TestCase):
 
         fut = asyncio.Future()
         fut.set_result(UNKNOWN)
-        out = Output.from_input({ "foo": "foo", "bar": UNKNOWN, "baz": fut})
+        out = Output.from_input({"foo": "foo", "bar": UNKNOWN, "baz": fut})
 
         self.assertFalse(await out.is_known())
 
@@ -855,7 +855,7 @@ class NextSerializationTests(unittest.TestCase):
         self.assertFalse(await r4.is_known())
         self.assertEqual(await r4.future(with_unknowns=True), UNKNOWN)
 
-        out = Output.from_input([ "foo", UNKNOWN ])
+        out = Output.from_input(["foo", UNKNOWN])
 
         r5 = out[0]
         self.assertTrue(await r5.is_known())
@@ -866,7 +866,7 @@ class NextSerializationTests(unittest.TestCase):
         self.assertEqual(await r6.future(with_unknowns=True), UNKNOWN)
 
         out = Output.all(Output.from_input("foo"), Output.from_input(UNKNOWN),
-            Output.from_input([ Output.from_input(UNKNOWN), Output.from_input("bar") ]))
+                         Output.from_input([Output.from_input(UNKNOWN), Output.from_input("bar")]))
 
         self.assertFalse(await out.is_known())
 
@@ -889,7 +889,6 @@ class NextSerializationTests(unittest.TestCase):
         self.assertTrue(await r11.is_known())
         self.assertEqual(await r11.future(with_unknowns=True), "bar")
 
-
     @async_test
     async def test_output_coros(self):
         # Ensure that Outputs function properly when the input value and is_known are coroutines. If the implementation
@@ -898,6 +897,7 @@ class NextSerializationTests(unittest.TestCase):
         async def value():
             await asyncio.sleep(0)
             return 42
+
         async def is_known():
             await asyncio.sleep(0)
             return True
@@ -922,7 +922,7 @@ class DeserializationTests(unittest.TestCase):
         self.assertIsNotNone(error)
 
     def test_secret_push_up(self):
-        secret_value = {rpc._special_sig_key: rpc._special_secret_sig, "value": "a secret value" }
+        secret_value = {rpc._special_sig_key: rpc._special_secret_sig, "value": "a secret value"}
         all_props = struct_pb2.Struct()
         all_props["regular"] = "a normal value"
         all_props["list"] = ["a normal value", "another value", secret_value]
@@ -930,7 +930,6 @@ class DeserializationTests(unittest.TestCase):
         all_props["mapWithList"] = {"regular": "a normal value", "list": ["a normal value", secret_value]}
         all_props["listWithMap"] = [{"regular": "a normal value", "secret": secret_value}]
 
-
         val = rpc.deserialize_properties(all_props)
         self.assertEqual(all_props["regular"], val["regular"])
 
@@ -971,15 +970,17 @@ class DeserializationTests(unittest.TestCase):
             "__provider": "serialized_dynamic_provider",
         }, val)
 
+
 @input_type
 class FooArgs:
     first_arg: Input[str] = pulumi.property("firstArg")
     second_arg: Optional[Input[float]] = pulumi.property("secondArg")
 
-    def __init__(self, first_arg: Input[str], second_arg: Optional[Input[float]]=None):
+    def __init__(self, first_arg: Input[str], second_arg: Optional[Input[float]] = None):
         pulumi.set(self, "first_arg", first_arg)
         pulumi.set(self, "second_arg", second_arg)
 
+
 @input_type
 class ListDictInputArgs:
     a: List[Input[str]]
@@ -1043,3 +1044,39 @@ class InputTypeSerializationTests(unittest.TestCase):
                 "foo_baz": "world",
             },
         }, prop)
+
+
+class StrEnum(str, Enum):
+    ONE = "one"
+    ZERO = "zero"
+
+
+class IntEnum(int, Enum):
+    ONE = 1
+    ZERO = 0
+
+
+class FloatEnum(float, Enum):
+    ONE = 1.0
+    ZERO_POINT_ONE = 0.1
+
+
+class EnumSerializationTests(unittest.TestCase):
+    @async_test
+    async def test_string_enum(self):
+        one = StrEnum.ONE
+        prop = await rpc.serialize_property(one, [])
+        self.assertEqual(StrEnum.ONE, prop)
+
+    @async_test
+    async def test_int_enum(self):
+        one = IntEnum.ONE
+        prop = await rpc.serialize_property(one, [])
+        self.assertEqual(IntEnum.ONE, prop)
+
+    @async_test
+    async def test_float_enum(self):
+        one = FloatEnum.ZERO_POINT_ONE
+        prop = await rpc.serialize_property(one, [])
+        self.assertEqual(FloatEnum.ZERO_POINT_ONE, prop)
+

sdk/python/lib/test/test_translate_output_properties.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import unittest
+from enum import Enum
 from typing import Any, Dict, List, NamedTuple, Mapping, Optional, Sequence
 
 from pulumi.runtime import rpc
@@ -598,7 +599,43 @@ class OutputTypeWithAny(dict):
     value_str: Any
 
 
+class ContainerColor(str, Enum):
+    RED = "red"
+    BLUE = "blue"
+
+
+class ContainerSize(int, Enum):
+    FOUR_INCH = 4
+    SIX_INCH = 6
+
+
+class ContainerBrightness(float, Enum):
+    ZERO_POINT_ONE = 0.1
+    ONE_POINT_ZERO = 1.0
+
+
 class TranslateOutputPropertiesTests(unittest.TestCase):
+    def test_str_enum(self):
+        result = rpc.translate_output_properties("red", translate_output_property, ContainerColor)
+        self.assertIsInstance(result, ContainerColor)
+        self.assertIsInstance(result, Enum)
+        self.assertEqual(result, "red")
+        self.assertEqual(result, ContainerColor.RED)
+
+    def test_int_enum(self):
+        result = rpc.translate_output_properties(4, translate_output_property, ContainerSize)
+        self.assertIsInstance(result, ContainerSize)
+        self.assertIsInstance(result, Enum)
+        self.assertEqual(result, 4)
+        self.assertEqual(result, ContainerSize.FOUR_INCH)
+
+    def test_float_enum(self):
+        result = rpc.translate_output_properties(0.1, translate_output_property, ContainerBrightness)
+        self.assertIsInstance(result, ContainerBrightness)
+        self.assertIsInstance(result, Enum)
+        self.assertEqual(result, 0.1)
+        self.assertEqual(result, ContainerBrightness.ZERO_POINT_ONE)
+
     def test_translate(self):
         output = {
             "firstArg": "hello",

tests/integration/enums/python/.gitignore

@@ -0,0 +1,2 @@
+*.pyc
+venv/

tests/integration/enums/python/Pulumi.yaml

@@ -0,0 +1,3 @@
+name: enum_outputs_python
+runtime: python
+description: Enum outputs in python

tests/integration/enums/python/__main__.py

@@ -0,0 +1,43 @@
+from pulumi import Input, Output, export
+from pulumi.dynamic import Resource, ResourceProvider, CreateResult
+from enum import Enum
+from typing import Optional, Union
+
+
+class RubberTreeVariety(str, Enum):
+    BURGUNDY = "Burgundy"
+    RUBY = "Ruby"
+    TINEKE = "Tineke"
+
+
+class Farm(str, Enum):
+    PLANTS_R_US = "Plants'R'Us"
+    PULUMI_PLANTERS_INC = "Pulumi Planters Inc."
+
+
+current_id = 0
+
+
+class PlantProvider(ResourceProvider):
+    def create(self, inputs):
+        global current_id
+        current_id += 1
+        return CreateResult(str(current_id), inputs)
+
+
+class Tree(Resource):
+    type: Output[RubberTreeVariety]
+    farm: Optional[Output[str]]
+
+    def __init__(self, name: str, type: Input[RubberTreeVariety], farm: Optional[Input[str]]):
+        self.type = type
+        self.farm = farm
+        super().__init__(PlantProvider(), name, {"type": type, "farm": farm})
+
+
+# Create a resource with input object.
+tree = Tree("myTree", type=RubberTreeVariety.BURGUNDY, farm=Farm.PULUMI_PLANTERS_INC)
+
+export("myTreeType", tree.type)
+export("myTreeFarmChanged", tree.farm.apply(lambda x: x + "foo"))
+export("mySentence", Output.all(tree.type, tree.farm).apply(lambda args: f"My {args[0]} Rubber tree is from {args[1]}"))

tests/integration/enums/python/requirements.txt

@@ -0,0 +1 @@
+pulumi>=2.0.0,<3.0.0

tests/integration/integration_python_test.go

@@ -349,6 +349,23 @@ func TestLargeResourcePython(t *testing.T) {
 	})
 }
 
+// Test enum outputs
+func TestEnumOutputsPython(t *testing.T) {
+	integration.ProgramTest(t, &integration.ProgramTestOptions{
+		Dependencies: []string{
+			filepath.Join("..", "..", "sdk", "python", "env", "src"),
+		},
+		Dir: filepath.Join("enums", "python"),
+		ExtraRuntimeValidation: func(t *testing.T, stack integration.RuntimeValidationStackInfo) {
+			assert.NotNil(t, stack.Outputs)
+
+			assert.Equal(t, "Burgundy", stack.Outputs["myTreeType"])
+			assert.Equal(t, "Pulumi Planters Inc.foo", stack.Outputs["myTreeFarmChanged"])
+			assert.Equal(t, "My Burgundy Rubber tree is from Pulumi Planters Inc.", stack.Outputs["mySentence"])
+		},
+	})
+}
+
 // Test to ensure Pylint is clean.
 func TestPythonPylint(t *testing.T) {
 	t.Skip("Temporarily skipping test - pulumi/pulumi#4849")