diff --git a/src/google/protobuf/compiler/python/generator.cc b/src/google/protobuf/compiler/python/generator.cc index ebfe4ce26d..5d2fd4aeb1 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 0d1bf5929e..bc98b1f57e 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* seen_aliases) const { + const FileDescriptor& desc, absl::flat_hash_set* 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 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 fca8654afd..445d54efa5 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* seen_aliases) const; + void PrintImportForDescriptor(const FileDescriptor& desc, + absl::flat_hash_set* seen_aliases, + bool* has_importlib) const; template void Annotate(const std::string& label, const DescriptorT* descriptor) const; void PrintImports() const;