From 4cc258a6a1cb6d3e526e3ca2f6ab7748ee84e2e3 Mon Sep 17 00:00:00 2001
From: Jie Luo <jieluo@google.com>
Date: Thu, 20 Jul 2023 13:00:35 -0700
Subject: [PATCH] Add protobuf python keyword support in path for stub
 generator. Also added tests for keywords in path

PiperOrigin-RevId: 549715333
---
 .../protobuf/compiler/python/generator.cc     |  6 ++-
 .../protobuf/compiler/python/pyi_generator.cc | 47 +++++++++++--------
 .../protobuf/compiler/python/pyi_generator.h  |  6 +--
 3 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/src/google/protobuf/compiler/python/generator.cc b/src/google/protobuf/compiler/python/generator.cc
index ebfe4ce26d3e..5d2fd4aeb156 100644
--- a/src/google/protobuf/compiler/python/generator.cc
+++ b/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.
-      printer_->Print("import importlib\n");
+      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 {
diff --git a/src/google/protobuf/compiler/python/pyi_generator.cc b/src/google/protobuf/compiler/python/pyi_generator.cc
index 0d1bf5929e5a..bc98b1f57e0b 100644
--- a/src/google/protobuf/compiler/python/pyi_generator.cc
+++ b/src/google/protobuf/compiler/python/pyi_generator.cc
@@ -151,41 +151,50 @@ 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 import_statement;
-  if (last_dot_pos == std::string::npos) {
-    import_statement = absl::StrCat("import ", module_name);
-  } else {
-    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);
+  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");
   }
-  printer_->Print("$statement$ as $alias$\n", "statement",
-                  import_statement, "alias", alias);
-  import_map_[filename] = alias;
-  seen_aliases->insert(alias);
+  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);
+    } else {
+      import_statement =
+          absl::StrCat("from ", module_name.substr(0, last_dot_pos), " import ",
+                       module_name.substr(last_dot_pos + 1));
+    }
+    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);
     }
   }
 
diff --git a/src/google/protobuf/compiler/python/pyi_generator.h b/src/google/protobuf/compiler/python/pyi_generator.h
index fca8654afd20..445d54efa569 100644
--- a/src/google/protobuf/compiler/python/pyi_generator.h
+++ b/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;