Add protobuf python keyword support in path for stub generator. Also added tests for

keywords in path

PiperOrigin-RevId: 549715333

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

@@ -382,6 +382,7 @@ void Generator::PrintTopBoilerplate() const {
 
 // Prints Python imports for all modules imported by |file|.
 void Generator::PrintImports() const {
+  bool has_importlib = false;
   for (int i = 0; i < file_->dependency_count(); ++i) {
     absl::string_view filename = file_->dependency(i)->name();
 
@@ -396,7 +397,10 @@ void Generator::PrintImports() const {
       // module name and import it using importlib. Otherwise the usual kind of
       // import statement would result in a syntax error from the presence of
       // the keyword.
+      if (has_importlib == false) {
         printer_->Print("import importlib\n");
+        has_importlib = true;
+      }
       printer_->Print("$alias$ = importlib.import_module('$name$')\n", "alias",
                       module_alias, "name", module_name);
     } else {

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

@@ -151,12 +151,25 @@ void CheckImportModules(const Descriptor* descriptor,
 }
 
 void PyiGenerator::PrintImportForDescriptor(
-    const FileDescriptor& desc,
-    absl::flat_hash_set<std::string>* seen_aliases) const {
+    const FileDescriptor& desc, absl::flat_hash_set<std::string>* seen_aliases,
+    bool* has_importlib) const {
   const std::string& filename = desc.name();
   std::string module_name_owned = StrippedModuleName(filename);
   absl::string_view module_name(module_name_owned);
   size_t last_dot_pos = module_name.rfind('.');
+  std::string alias = absl::StrCat("_", module_name.substr(last_dot_pos + 1));
+  // Generate a unique alias by adding _1 suffixes until we get an unused alias.
+  while (seen_aliases->find(alias) != seen_aliases->end()) {
+    absl::StrAppend(&alias, "_1");
+  }
+  if (ContainsPythonKeyword(module_name)) {
+    if (*has_importlib == false) {
+      printer_->Print("import importlib\n");
+      *has_importlib = true;
+    }
+    printer_->Print("$alias$ = importlib.import_module('$name$')\n", "alias",
+                    alias, "name", module_name);
+  } else {
     std::string import_statement;
     if (last_dot_pos == std::string::npos) {
       import_statement = absl::StrCat("import ", module_name);
@@ -164,28 +177,24 @@ void PyiGenerator::PrintImportForDescriptor(
       import_statement =
           absl::StrCat("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);
     }
-  std::string alias = absl::StrCat("_", module_name);
-  // Generate a unique alias by adding _1 suffixes until we get an unused alias.
-  while (seen_aliases->find(alias) != seen_aliases->end()) {
-    absl::StrAppend(&alias, "_1");
-  }
-  printer_->Print("$statement$ as $alias$\n", "statement",
-                  import_statement, "alias", alias);
+    printer_->Print("$statement$ as $alias$\n", "statement", import_statement,
+                    "alias", alias);
     import_map_[filename] = alias;
     seen_aliases->insert(alias);
+  }
 }
 
 void PyiGenerator::PrintImports() const {
   // Prints imported dependent _pb2 files.
   absl::flat_hash_set<std::string> seen_aliases;
+  bool has_importlib = false;
   for (int i = 0; i < file_->dependency_count(); ++i) {
     const FileDescriptor* dep = file_->dependency(i);
-    PrintImportForDescriptor(*dep, &seen_aliases);
+    PrintImportForDescriptor(*dep, &seen_aliases, &has_importlib);
     for (int j = 0; j < dep->public_dependency_count(); ++j) {
-      PrintImportForDescriptor(
-          *dep->public_dependency(j), &seen_aliases);
+      PrintImportForDescriptor(*dep->public_dependency(j), &seen_aliases,
+                               &has_importlib);
     }
   }
 

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

@@ -77,9 +77,9 @@ class PROTOC_EXPORT PyiGenerator : public google::protobuf::compiler::CodeGenera
                 std::string* error) const override;
 
  private:
-  void PrintImportForDescriptor(
-      const FileDescriptor& desc,
-      absl::flat_hash_set<std::string>* seen_aliases) const;
+  void PrintImportForDescriptor(const FileDescriptor& desc,
+                                absl::flat_hash_set<std::string>* seen_aliases,
+                                bool* has_importlib) const;
   template <typename DescriptorT>
   void Annotate(const std::string& label, const DescriptorT* descriptor) const;
   void PrintImports() const;