Add type annotation for enum

PiperOrigin-RevId: 486439287

src/google/protobuf/compiler/python/pyi_generator.cc

@@ -304,17 +304,26 @@ void PyiGenerator::PrintEnum(const EnumDescriptor& enum_descriptor) const {
       "    __slots__ = []\n",
       "enum_name", enum_name);
   Annotate("enum_name", &enum_descriptor);
+  printer_->Indent();
+  PrintEnumValues(enum_descriptor, /* is_classvar = */ true);
+  printer_->Outdent();
 }
 
-void PyiGenerator::PrintEnumValues(
-    const EnumDescriptor& enum_descriptor) const {
+void PyiGenerator::PrintEnumValues(const EnumDescriptor& enum_descriptor,
+                                   bool is_classvar) const {
   // enum values
   std::string module_enum_name = ModuleLevelName(enum_descriptor);
   for (int j = 0; j < enum_descriptor.value_count(); ++j) {
     const EnumValueDescriptor* value_descriptor = enum_descriptor.value(j);
-    printer_->Print("$name$: $module_enum_name$\n",
-                    "name", value_descriptor->name(),
-                    "module_enum_name", module_enum_name);
+    if (is_classvar) {
+      printer_->Print("$name$: _ClassVar[$module_enum_name$]\n", "name",
+                      value_descriptor->name(), "module_enum_name",
+                      module_enum_name);
+    } else {
+      printer_->Print("$name$: $module_enum_name$\n", "name",
+                      value_descriptor->name(), "module_enum_name",
+                      module_enum_name);
+    }
     Annotate("name", value_descriptor);
   }
 }
@@ -398,7 +407,6 @@ void PyiGenerator::PrintMessage(
                   "class_name", class_name, "extra_base", extra_base);
   Annotate("class_name", &message_descriptor);
   printer_->Indent();
-  printer_->Indent();
 
   // Prints slots
   printer_->Print("__slots__ = [", "class_name", class_name);
@@ -538,8 +546,6 @@ void PyiGenerator::PrintMessage(
     printer_->Print(", **kwargs");
   }
   printer_->Print(") -> None: ...\n");
-
-  printer_->Outdent();
   printer_->Outdent();
 }
 
@@ -597,8 +603,10 @@ bool PyiGenerator::Generate(const FileDescriptor* file,
   GeneratedCodeInfo annotations;
   io::AnnotationProtoCollector<GeneratedCodeInfo> annotation_collector(
       &annotations);
-  io::Printer printer(output.get(), '$',
-                      annotate_code ? &annotation_collector : nullptr);
+  io::Printer::Options printer_opt(
+      '$', annotate_code ? &annotation_collector : nullptr);
+  printer_opt.spaces_per_indent = 4;
+  io::Printer printer(output.get(), printer_opt);
   printer_ = &printer;
 
   PrintImports();

src/google/protobuf/compiler/python/pyi_generator.h

@@ -85,7 +85,8 @@ class PROTOC_EXPORT PyiGenerator : public google::protobuf::compiler::CodeGenera
   void PrintImports() const;
   void PrintTopLevelEnums() const;
   void PrintEnum(const EnumDescriptor& enum_descriptor) const;
-  void PrintEnumValues(const EnumDescriptor& enum_descriptor) const;
+  void PrintEnumValues(const EnumDescriptor& enum_descriptor,
+                       bool is_classvar = false) const;
   template <typename DescriptorT>
   void PrintExtensions(const DescriptorT& descriptor) const;
   void PrintMessages() const;