Python: Fix enums not being casted to their types in generator

1 files changed, 10 insertions, 0 deletions

diff --git a/python/generator.cpp b/python/generator.cpp
index 41fbb159..6dd983b4 100644
--- a/python/generator.cpp
+++ b/python/generator.cpp
@@ -470,6 +470,9 @@ int main(int argc, char* argv[])
 		                    && (i.second->GetChildType()->GetChildType()->IsSigned());
 		// Pointer returns will be automatically wrapped to return None on null pointer
 		bool pointerResult = (i.second->GetChildType()->GetClass() == PointerTypeClass);
+		// Enum returns will automatically cast to the enum type
+		bool enumResult = (i.second->GetChildType()->GetClass() == NamedTypeReferenceClass
+		                   && i.second->GetChildType()->GetNamedTypeReference()->GetTypeReferenceClass() == EnumNamedTypeClass);
 
 		// From python -> C python3 requires str -> str.encode('charmap')
 		bool swizzleArgs = true;
@@ -617,6 +620,13 @@ int main(int argc, char* argv[])
 			fprintf(out, "\t\treturn None\n");
 			fprintf(out, "\treturn result\n");
 		}
+		else if (enumResult)
+		{
+			// Emit wrapper to cast result to enum type
+			fprintf(out, "\treturn ");
+			OutputSwizzledType(out, i.second->GetChildType().GetValue());
+			fprintf(out, "(%s)", stringArgFuncCall.c_str());
+		}
 		else
 		{
 			fprintf(out, "\treturn ");