Update protobuf to 3.21.2

a628f0376085fcf46dc6d24629f2a7dacb91ae79

6 files changed, 180 insertions, 91 deletions

diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_generator.cc b/contrib/libs/protoc/src/google/protobuf/compiler/python/generator.cc
index 7c723febb38..d1e4de467b9 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_generator.cc
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/generator.cc
@@ -42,7 +42,7 @@
 // performance-minded Python code leverage the fast C++ implementation
 // directly.
 
-#include <google/protobuf/compiler/python/python_generator.h>
+#include <google/protobuf/compiler/python/generator.h>
 
 #include <algorithm>
 #include <limits>
@@ -54,15 +54,15 @@
 
 #include <google/protobuf/stubs/logging.h>
 #include <google/protobuf/stubs/common.h>
-#include <google/protobuf/compiler/python/python_helpers.h>
-#include <google/protobuf/compiler/python/python_pyi_generator.h>
-#include <google/protobuf/io/printer.h>
-#include <google/protobuf/io/zero_copy_stream.h>
-#include <google/protobuf/descriptor.h>
 #include <google/protobuf/stubs/strutil.h>
 #include <google/protobuf/stubs/stringprintf.h>
 #include <google/protobuf/stubs/substitute.h>
+#include <google/protobuf/compiler/python/helpers.h>
+#include <google/protobuf/compiler/python/pyi_generator.h>
+#include <google/protobuf/descriptor.h>
 #include <google/protobuf/descriptor.pb.h>
+#include <google/protobuf/io/printer.h>
+#include <google/protobuf/io/zero_copy_stream.h>
 
 namespace google {
 namespace protobuf {
@@ -91,11 +91,6 @@ TProtoStringType ModuleAlias(const TProtoStringType& filename) {
 const char kDescriptorKey[] = "DESCRIPTOR";
 
 
-// Does the file have top-level enums?
-inline bool HasTopLevelEnums(const FileDescriptor* file) {
-  return file->enum_type_count() > 0;
-}
-
 // file output by this generator.
 void PrintTopBoilerplate(io::Printer* printer, const FileDescriptor* file,
                          bool descriptor_proto) {
diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_generator.h b/contrib/libs/protoc/src/google/protobuf/compiler/python/generator.h
index 359d802aeea..8111bf733fc 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_generator.h
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/generator.h
@@ -37,8 +37,8 @@
 
 #include <string>
 
-#include <google/protobuf/compiler/code_generator.h>
 #include <google/protobuf/stubs/mutex.h>
+#include <google/protobuf/compiler/code_generator.h>
 
 // Must be included last.
 #include <google/protobuf/port_def.inc>
diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_helpers.cc b/contrib/libs/protoc/src/google/protobuf/compiler/python/helpers.cc
index 7a9721c02ca..ee7461bc83c 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_helpers.cc
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/helpers.cc
@@ -28,13 +28,13 @@
 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-#include <google/protobuf/compiler/python/python_helpers.h>
+#include <google/protobuf/compiler/python/helpers.h>
 
 #include <algorithm>
 
+#include <google/protobuf/stubs/strutil.h>
 #include <google/protobuf/compiler/code_generator.h>
 #include <google/protobuf/descriptor.h>
-#include <google/protobuf/stubs/strutil.h>
 #include <google/protobuf/descriptor.pb.h>
 
 namespace google {
@@ -69,7 +69,7 @@ const char* const kKeywords[] = {
     "del",    "elif",     "else",  "except", "finally",  "for",
     "from",   "global",   "if",    "import", "in",       "is",
     "lambda", "nonlocal", "not",   "or",     "pass",     "raise",
-    "return", "try",      "while", "with",   "yield",    "print",
+    "return", "try",      "while", "with",   "yield",
 };
 const char* const* kKeywordsEnd =
     kKeywords + (sizeof(kKeywords) / sizeof(kKeywords[0]));
diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_helpers.h b/contrib/libs/protoc/src/google/protobuf/compiler/python/helpers.h
index 8745698acf6..8745698acf6 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_helpers.h
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/helpers.h
diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_pyi_generator.cc b/contrib/libs/protoc/src/google/protobuf/compiler/python/pyi_generator.cc
index 89cfbf86fe6..242cdd71f2c 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_pyi_generator.cc
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/pyi_generator.cc
@@ -28,16 +28,16 @@
 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-#include <google/protobuf/compiler/python/python_pyi_generator.h>
+#include <google/protobuf/compiler/python/pyi_generator.h>
 
 #include <string>
 
-#include <google/protobuf/compiler/python/python_helpers.h>
-#include <google/protobuf/io/printer.h>
-#include <google/protobuf/io/zero_copy_stream.h>
-#include <google/protobuf/descriptor.h>
 #include <google/protobuf/stubs/strutil.h>
+#include <google/protobuf/compiler/python/helpers.h>
+#include <google/protobuf/descriptor.h>
 #include <google/protobuf/descriptor.pb.h>
+#include <google/protobuf/io/printer.h>
+#include <google/protobuf/io/zero_copy_stream.h>
 
 namespace google {
 namespace protobuf {
@@ -64,12 +64,21 @@ void PyiGenerator::PrintItemMap(
 }
 
 template <typename DescriptorT>
-TProtoStringType PyiGenerator::ModuleLevelName(const DescriptorT& descriptor) const {
+TProtoStringType PyiGenerator::ModuleLevelName(
+    const DescriptorT& descriptor,
+    const std::map<TProtoStringType, TProtoStringType>& import_map) const {
   TProtoStringType name = NamePrefixedWithNestedTypes(descriptor, ".");
   if (descriptor.file() != file_) {
-    TProtoStringType module_name = ModuleName(descriptor.file()->name());
-    std::vector<TProtoStringType> tokens = Split(module_name, ".");
-    name = "_" + tokens.back() + "." + name;
+    TProtoStringType module_alias;
+    TProtoStringType filename = descriptor.file()->name();
+    if (import_map.find(filename) == import_map.end()) {
+      TProtoStringType module_name = ModuleName(descriptor.file()->name());
+      std::vector<TProtoStringType> tokens = Split(module_name, ".");
+      module_alias = "_" + tokens.back();
+    } else {
+      module_alias = import_map.at(filename);
+    }
+    name = module_alias + "." + name;
   }
   return name;
 }
@@ -82,9 +91,22 @@ struct ImportModules {
   bool has_extendable = false;  // _python_message
   bool has_mapping = false;     // typing.Mapping
   bool has_optional = false;    // typing.Optional
-  bool has_union = false;       // typing.Uion
+  bool has_union = false;       // typing.Union
+  bool has_well_known_type = false;
 };
 
+// Checks whether a descriptor name matches a well-known type.
+bool IsWellKnownType(const TProtoStringType& name) {
+  // LINT.IfChange(wktbases)
+  return (name == "google.protobuf.Any" ||
+          name == "google.protobuf.Duration" ||
+          name == "google.protobuf.FieldMask" ||
+          name == "google.protobuf.ListValue" ||
+          name == "google.protobuf.Struct" ||
+          name == "google.protobuf.Timestamp");
+  // LINT.ThenChange(//depot/google3/net/proto2/python/internal/well_known_types.py:wktbases)
+}
+
 // Checks what modules should be imported for this message
 // descriptor.
 void CheckImportModules(const Descriptor* descriptor,
@@ -95,6 +117,9 @@ void CheckImportModules(const Descriptor* descriptor,
   if (descriptor->enum_type_count() > 0) {
     import_modules->has_enums = true;
   }
+  if (IsWellKnownType(descriptor->full_name())) {
+    import_modules->has_well_known_type = true;
+  }
   for (int i = 0; i < descriptor->field_count(); ++i) {
     const FieldDescriptor* field = descriptor->field(i);
     if (IsPythonKeyword(field->name())) {
@@ -129,23 +154,44 @@ void CheckImportModules(const Descriptor* descriptor,
   }
 }
 
+void PyiGenerator::PrintImportForDescriptor(
+    const FileDescriptor& desc,
+    std::map<TProtoStringType, TProtoStringType>* import_map,
+    std::set<TProtoStringType>* seen_aliases) const {
+  const TProtoStringType& filename = desc.name();
+  TProtoStringType module_name = StrippedModuleName(filename);
+  size_t last_dot_pos = module_name.rfind('.');
+  TProtoStringType import_statement;
+  if (last_dot_pos == TProtoStringType::npos) {
+    import_statement = "import " + module_name;
+  } else {
+    import_statement = "from " + module_name.substr(0, last_dot_pos) +
+                       " import " + module_name.substr(last_dot_pos + 1);
+    module_name = module_name.substr(last_dot_pos + 1);
+  }
+  TProtoStringType alias = "_" + module_name;
+  // Generate a unique alias by adding _1 suffixes until we get an unused alias.
+  while (seen_aliases->find(alias) != seen_aliases->end()) {
+    alias = alias + "_1";
+  }
+  printer_->Print("$statement$ as $alias$\n", "statement",
+                  import_statement, "alias", alias);
+  (*import_map)[filename] = alias;
+  seen_aliases->insert(alias);
+}
+
 void PyiGenerator::PrintImports(
-    std::map<TProtoStringType, TProtoStringType>* item_map) const {
+    std::map<TProtoStringType, TProtoStringType>* item_map,
+    std::map<TProtoStringType, TProtoStringType>* import_map) const {
   // Prints imported dependent _pb2 files.
+  std::set<TProtoStringType> seen_aliases;
   for (int i = 0; i < file_->dependency_count(); ++i) {
-    const TProtoStringType& filename = file_->dependency(i)->name();
-    TProtoStringType module_name = StrippedModuleName(filename);
-    size_t last_dot_pos = module_name.rfind('.');
-    TProtoStringType import_statement;
-    if (last_dot_pos == TProtoStringType::npos) {
-      import_statement = "import " + module_name;
-    } else {
-      import_statement = "from " + module_name.substr(0, last_dot_pos) +
-                         " import " + module_name.substr(last_dot_pos + 1);
-      module_name = module_name.substr(last_dot_pos + 1);
+    const FileDescriptor* dep = file_->dependency(i);
+    PrintImportForDescriptor(*dep, import_map, &seen_aliases);
+    for (int j = 0; j < dep->public_dependency_count(); ++j) {
+      PrintImportForDescriptor(
+          *dep->public_dependency(j), import_map, &seen_aliases);
     }
-    printer_->Print("$statement$ as _$module_name$\n", "statement",
-                    import_statement, "module_name", module_name);
   }
 
   // Checks what modules should be imported.
@@ -177,6 +223,11 @@ void PyiGenerator::PrintImports(
         "from google.protobuf.internal import python_message"
         " as _python_message\n");
   }
+  if (import_modules.has_well_known_type) {
+    printer_->Print(
+        "from google.protobuf.internal import well_known_types"
+        " as _well_known_types\n");
+  }
   printer_->Print(
       "from google.protobuf import"
       " descriptor as _descriptor\n");
@@ -190,21 +241,18 @@ void PyiGenerator::PrintImports(
         " _service\n");
   }
   printer_->Print("from typing import ");
-  printer_->Print("ClassVar");
+  printer_->Print("ClassVar as _ClassVar");
   if (import_modules.has_iterable) {
-    printer_->Print(", Iterable");
+    printer_->Print(", Iterable as _Iterable");
   }
   if (import_modules.has_mapping) {
-    printer_->Print(", Mapping");
+    printer_->Print(", Mapping as _Mapping");
   }
   if (import_modules.has_optional) {
-    printer_->Print(", Optional");
-  }
-  if (file_->service_count() > 0) {
-    printer_->Print(", Text");
+    printer_->Print(", Optional as _Optional");
   }
   if (import_modules.has_union) {
-    printer_->Print(", Union");
+    printer_->Print(", Union as _Union");
   }
   printer_->Print("\n\n");
 
@@ -229,7 +277,7 @@ void PyiGenerator::PrintImports(
       const EnumDescriptor* enum_descriptor = public_dep->enum_type(i);
       for (int j = 0; j < enum_descriptor->value_count(); ++j) {
         (*item_map)[enum_descriptor->value(j)->name()] =
-            ModuleLevelName(*enum_descriptor);
+            ModuleLevelName(*enum_descriptor, *import_map);
       }
     }
     // Top level extensions for public imports
@@ -248,9 +296,10 @@ void PyiGenerator::PrintEnum(const EnumDescriptor& enum_descriptor) const {
 // Adds enum value to item map which will be ordered and printed later.
 void PyiGenerator::AddEnumValue(
     const EnumDescriptor& enum_descriptor,
-    std::map<TProtoStringType, TProtoStringType>* item_map) const {
+    std::map<TProtoStringType, TProtoStringType>* item_map,
+    const std::map<TProtoStringType, TProtoStringType>& import_map) const {
   // enum values
-  TProtoStringType module_enum_name = ModuleLevelName(enum_descriptor);
+  TProtoStringType module_enum_name = ModuleLevelName(enum_descriptor, import_map);
   for (int j = 0; j < enum_descriptor.value_count(); ++j) {
     const EnumValueDescriptor* value_descriptor = enum_descriptor.value(j);
     (*item_map)[value_descriptor->name()] = module_enum_name;
@@ -275,13 +324,15 @@ void PyiGenerator::AddExtensions(
     const FieldDescriptor* extension_field = descriptor.extension(i);
     TProtoStringType constant_name = extension_field->name() + "_FIELD_NUMBER";
     ToUpper(&constant_name);
-    (*item_map)[constant_name] = "ClassVar[int]";
+    (*item_map)[constant_name] = "_ClassVar[int]";
     (*item_map)[extension_field->name()] = "_descriptor.FieldDescriptor";
   }
 }
 
 // Returns the string format of a field's cpp_type
-TProtoStringType PyiGenerator::GetFieldType(const FieldDescriptor& field_des) const {
+TProtoStringType PyiGenerator::GetFieldType(
+    const FieldDescriptor& field_des, const Descriptor& containing_des,
+    const std::map<TProtoStringType, TProtoStringType>& import_map) const {
   switch (field_des.cpp_type()) {
     case FieldDescriptor::CPPTYPE_INT32:
     case FieldDescriptor::CPPTYPE_UINT32:
@@ -294,29 +345,48 @@ TProtoStringType PyiGenerator::GetFieldType(const FieldDescriptor& field_des) co
     case FieldDescriptor::CPPTYPE_BOOL:
       return "bool";
     case FieldDescriptor::CPPTYPE_ENUM:
-      return ModuleLevelName(*field_des.enum_type());
+      return ModuleLevelName(*field_des.enum_type(), import_map);
     case FieldDescriptor::CPPTYPE_STRING:
       if (field_des.type() == FieldDescriptor::TYPE_STRING) {
         return "str";
       } else {
         return "bytes";
       }
-    case FieldDescriptor::CPPTYPE_MESSAGE:
-      return ModuleLevelName(*field_des.message_type());
+    case FieldDescriptor::CPPTYPE_MESSAGE: {
+      // If the field is inside a nested message and the nested message has the
+      // same name as a top-level message, then we need to prefix the field type
+      // with the module name for disambiguation.
+      TProtoStringType name = ModuleLevelName(*field_des.message_type(), import_map);
+      if ((containing_des.containing_type() != nullptr &&
+           name == containing_des.name())) {
+        TProtoStringType module = ModuleName(field_des.file()->name());
+        name = module + "." + name;
+      }
+      return name;
+    }
     default:
-      GOOGLE_LOG(FATAL) << "Unsuppoted field type.";
+      GOOGLE_LOG(FATAL) << "Unsupported field type.";
   }
   return "";
 }
 
-void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
-                                bool is_nested) const {
+void PyiGenerator::PrintMessage(
+    const Descriptor& message_descriptor, bool is_nested,
+    const std::map<TProtoStringType, TProtoStringType>& import_map) const {
   if (!is_nested) {
     printer_->Print("\n");
   }
   TProtoStringType class_name = message_descriptor.name();
-  printer_->Print("class $class_name$(_message.Message):\n", "class_name",
-                  class_name);
+  TProtoStringType extra_base;
+  // A well-known type needs to inherit from its corresponding base class in
+  // net/proto2/python/internal/well_known_types.
+  if (IsWellKnownType(message_descriptor.full_name())) {
+    extra_base = ", _well_known_types." + message_descriptor.name();
+  } else {
+    extra_base = "";
+  }
+  printer_->Print("class $class_name$(_message.Message$extra_base$):\n",
+                  "class_name", class_name, "extra_base", extra_base);
   printer_->Indent();
   printer_->Indent();
 
@@ -361,7 +431,7 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
   for (const auto& entry : nested_enums) {
     PrintEnum(*entry);
     // Adds enum value to item_map which will be ordered and printed later
-    AddEnumValue(*entry, &item_map);
+    AddEnumValue(*entry, &item_map, import_map);
   }
 
   // Prints nested messages
@@ -374,7 +444,7 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
             SortByName<Descriptor>());
 
   for (const auto& entry : nested_messages) {
-    PrintMessage(*entry, true);
+    PrintMessage(*entry, true, import_map);
   }
 
   // Adds extensions to item_map which will be ordered and printed later
@@ -384,7 +454,7 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
   for (int i = 0; i < message_descriptor.field_count(); ++i) {
     const FieldDescriptor& field_des = *message_descriptor.field(i);
     item_map[ToUpper(field_des.name()) + "_FIELD_NUMBER"] =
-        "ClassVar[int]";
+        "_ClassVar[int]";
     if (IsPythonKeyword(field_des.name())) {
       continue;
     }
@@ -395,16 +465,16 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
       field_type = (value_des->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE
                         ? "_containers.MessageMap["
                         : "_containers.ScalarMap[");
-      field_type += GetFieldType(*key_des);
+      field_type += GetFieldType(*key_des, message_descriptor, import_map);
       field_type += ", ";
-      field_type += GetFieldType(*value_des);
+      field_type += GetFieldType(*value_des, message_descriptor, import_map);
     } else {
       if (field_des.is_repeated()) {
         field_type = (field_des.cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE
                           ? "_containers.RepeatedCompositeFieldContainer["
                           : "_containers.RepeatedScalarFieldContainer[");
       }
-      field_type += GetFieldType(field_des);
+      field_type += GetFieldType(field_des, message_descriptor, import_map);
     }
 
     if (field_des.is_repeated()) {
@@ -437,26 +507,31 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
     printer_->Print(", $field_name$: ", "field_name", field_name);
     if (field_des->is_repeated() ||
         field_des->cpp_type() != FieldDescriptor::CPPTYPE_BOOL) {
-      printer_->Print("Optional[");
+      printer_->Print("_Optional[");
     }
     if (field_des->is_map()) {
       const Descriptor* map_entry = field_des->message_type();
-      printer_->Print("Mapping[$key_type$, $value_type$]", "key_type",
-                      GetFieldType(*map_entry->field(0)), "value_type",
-                      GetFieldType(*map_entry->field(1)));
+      printer_->Print(
+          "_Mapping[$key_type$, $value_type$]", "key_type",
+          GetFieldType(*map_entry->field(0), message_descriptor, import_map),
+          "value_type",
+          GetFieldType(*map_entry->field(1), message_descriptor, import_map));
     } else {
       if (field_des->is_repeated()) {
-        printer_->Print("Iterable[");
+        printer_->Print("_Iterable[");
       }
       if (field_des->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
-        printer_->Print("Union[$type_name$, Mapping]", "type_name",
-                        GetFieldType(*field_des));
+        printer_->Print(
+            "_Union[$type_name$, _Mapping]", "type_name",
+            GetFieldType(*field_des, message_descriptor, import_map));
       } else {
         if (field_des->cpp_type() == FieldDescriptor::CPPTYPE_ENUM) {
-          printer_->Print("Union[$type_name$, str]", "type_name",
-                          ModuleLevelName(*field_des->enum_type()));
+          printer_->Print("_Union[$type_name$, str]", "type_name",
+                          ModuleLevelName(*field_des->enum_type(), import_map));
         } else {
-          printer_->Print("$type_name$", "type_name", GetFieldType(*field_des));
+          printer_->Print(
+              "$type_name$", "type_name",
+              GetFieldType(*field_des, message_descriptor, import_map));
         }
       }
       if (field_des->is_repeated()) {
@@ -478,8 +553,9 @@ void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
   printer_->Outdent();
 }
 
-void PyiGenerator::PrintMessages() const {
-  // Order the descriptors by name to have same output with proto_to_pyi.py
+void PyiGenerator::PrintMessages(
+    const std::map<TProtoStringType, TProtoStringType>& import_map) const {
+  // Deterministically order the descriptors.
   std::vector<const Descriptor*> messages;
   messages.reserve(file_->message_type_count());
   for (int i = 0; i < file_->message_type_count(); ++i) {
@@ -488,7 +564,7 @@ void PyiGenerator::PrintMessages() const {
   std::sort(messages.begin(), messages.end(), SortByName<Descriptor>());
 
   for (const auto& entry : messages) {
-    PrintMessage(*entry, false);
+    PrintMessage(*entry, false, import_map);
   }
 }
 
@@ -517,9 +593,6 @@ bool PyiGenerator::Generate(const FileDescriptor* file,
   MutexLock lock(&mutex_);
   // Calculate file name.
   file_ = file;
-  // proto_to_pyi.py may set the output file name directly. To replace
-  // proto_to_pyi.py in google3, protoc also accept --pyi_out to set
-  // the output file name.
   TProtoStringType filename =
       parameter.empty() ? GetFileName(file, ".pyi") : parameter;
 
@@ -534,17 +607,22 @@ bool PyiGenerator::Generate(const FileDescriptor* file,
 
   // Adds "DESCRIPTOR" into item_map.
   item_map["DESCRIPTOR"] = "_descriptor.FileDescriptor";
-  PrintImports(&item_map);
+
+  // import_map will be a mapping from filename to module alias, e.g.
+  // "google3/foo/bar.py" -> "_bar"
+  std::map<TProtoStringType, TProtoStringType> import_map;
+
+  PrintImports(&item_map, &import_map);
   // Adds top level enum values to item_map.
   for (int i = 0; i < file_->enum_type_count(); ++i) {
-    AddEnumValue(*file_->enum_type(i), &item_map);
+    AddEnumValue(*file_->enum_type(i), &item_map, import_map);
   }
   // Adds top level extensions to item_map.
   AddExtensions(*file_, &item_map);
   // Prints item map
   PrintItemMap(item_map);
 
-  PrintMessages();
+  PrintMessages(import_map);
   PrintTopLevelEnums();
   if (HasGenericServices(file)) {
     PrintServices();
diff --git a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_pyi_generator.h b/contrib/libs/protoc/src/google/protobuf/compiler/python/pyi_generator.h
index d7f5437e491..c9bf2961665 100644
--- a/contrib/libs/protoc/src/google/protobuf/compiler/python/python_pyi_generator.h
+++ b/contrib/libs/protoc/src/google/protobuf/compiler/python/pyi_generator.h
@@ -36,10 +36,11 @@
 #define GOOGLE_PROTOBUF_COMPILER_PYTHON_PYI_GENERATOR_H__
 
 #include <map>
+#include <set>
 #include <string>
 
-#include <google/protobuf/compiler/code_generator.h>
 #include <google/protobuf/stubs/mutex.h>
+#include <google/protobuf/compiler/code_generator.h>
 
 // Must be included last.
 #include <google/protobuf/port_def.inc>
@@ -65,26 +66,41 @@ class PROTOC_EXPORT PyiGenerator : public google::protobuf::compiler::CodeGenera
   ~PyiGenerator() override;
 
   // CodeGenerator methods.
+  uint64_t GetSupportedFeatures() const override {
+    // Code generators must explicitly support proto3 optional.
+    return CodeGenerator::FEATURE_PROTO3_OPTIONAL;
+  }
   bool Generate(const FileDescriptor* file, const TProtoStringType& parameter,
                 GeneratorContext* generator_context,
                 TProtoStringType* error) const override;
 
  private:
-  void PrintImports(std::map<TProtoStringType, TProtoStringType>* item_map) const;
+  void PrintImportForDescriptor(const FileDescriptor& desc,
+                                std::map<TProtoStringType, TProtoStringType>* import_map,
+                                std::set<TProtoStringType>* seen_aliases) const;
+  void PrintImports(std::map<TProtoStringType, TProtoStringType>* item_map,
+                    std::map<TProtoStringType, TProtoStringType>* import_map) const;
   void PrintEnum(const EnumDescriptor& enum_descriptor) const;
   void AddEnumValue(const EnumDescriptor& enum_descriptor,
-                    std::map<TProtoStringType, TProtoStringType>* item_map) const;
+                    std::map<TProtoStringType, TProtoStringType>* item_map,
+                    const std::map<TProtoStringType, TProtoStringType>& import_map) const;
   void PrintTopLevelEnums() const;
   template <typename DescriptorT>
   void AddExtensions(const DescriptorT& descriptor,
                      std::map<TProtoStringType, TProtoStringType>* item_map) const;
-  void PrintMessages() const;
-  void PrintMessage(const Descriptor& message_descriptor, bool is_nested) const;
+  void PrintMessages(
+      const std::map<TProtoStringType, TProtoStringType>& import_map) const;
+  void PrintMessage(const Descriptor& message_descriptor, bool is_nested,
+                    const std::map<TProtoStringType, TProtoStringType>& import_map) const;
   void PrintServices() const;
   void PrintItemMap(const std::map<TProtoStringType, TProtoStringType>& item_map) const;
-  TProtoStringType GetFieldType(const FieldDescriptor& field_des) const;
+  TProtoStringType GetFieldType(
+      const FieldDescriptor& field_des, const Descriptor& containing_des,
+      const std::map<TProtoStringType, TProtoStringType>& import_map) const;
   template <typename DescriptorT>
-  TProtoStringType ModuleLevelName(const DescriptorT& descriptor) const;
+  TProtoStringType ModuleLevelName(
+      const DescriptorT& descriptor,
+      const std::map<TProtoStringType, TProtoStringType>& import_map) const;
 
   // Very coarse-grained lock to ensure that Generate() is reentrant.
   // Guards file_ and printer_.