diff options
Diffstat (limited to 'src/google/protobuf/compiler/python')
| -rw-r--r-- | src/google/protobuf/compiler/python/python_generator.cc | 60 | ||||
| -rw-r--r-- | src/google/protobuf/compiler/python/python_plugin_unittest.cc | 48 |
2 files changed, 92 insertions, 16 deletions
diff --git a/src/google/protobuf/compiler/python/python_generator.cc b/src/google/protobuf/compiler/python/python_generator.cc index cc54a8ab..d5468a0c 100644 --- a/src/google/protobuf/compiler/python/python_generator.cc +++ b/src/google/protobuf/compiler/python/python_generator.cc @@ -44,6 +44,7 @@ // performance-minded Python code leverage the fast C++ implementation // directly. +#include <algorithm> #include <google/protobuf/stubs/hash.h> #include <limits> #include <map> @@ -107,20 +108,25 @@ string ModuleAlias(const string& filename) { return module_name; } - -// Returns an import statement of form "from X.Y.Z import T" for the given -// .proto filename. -string ModuleImportStatement(const string& filename) { - string module_name = ModuleName(filename); - int last_dot_pos = module_name.rfind('.'); - if (last_dot_pos == string::npos) { - // NOTE(petya): this is not tested as it would require a protocol buffer - // outside of any package, and I don't think that is easily achievable. - return "import " + module_name; - } else { - return "from " + module_name.substr(0, last_dot_pos) + " import " + - module_name.substr(last_dot_pos + 1); +// Keywords reserved by the Python language. +const char* const kKeywords[] = { + "False", "None", "True", "and", "as", "assert", "break", + "class", "continue", "def", "del", "elif", "else", "except", + "finally", "for", "from", "global", "if", "import", "in", + "is", "lambda", "nonlocal", "not", "or", "pass", "raise", + "return", "try", "while", "with", "yield", +}; +const char* const* kKeywordsEnd = + kKeywords + (sizeof(kKeywords) / sizeof(kKeywords[0])); + +bool ContainsPythonKeyword(const string& module_name) { + vector<string> tokens = Split(module_name, "."); + for (int i = 0; i < tokens.size(); ++i) { + if (std::find(kKeywords, kKeywordsEnd, tokens[i]) != kKeywordsEnd) { + return true; + } } + return false; } @@ -359,10 +365,32 @@ bool Generator::Generate(const FileDescriptor* file, void Generator::PrintImports() const { for (int i = 0; i < file_->dependency_count(); ++i) { const string& filename = file_->dependency(i)->name(); - string import_statement = ModuleImportStatement(filename); + + string module_name = ModuleName(filename); string module_alias = ModuleAlias(filename); - printer_->Print("$statement$ as $alias$\n", "statement", - import_statement, "alias", module_alias); + if (ContainsPythonKeyword(module_name)) { + // If the module path contains a Python keyword, we have to quote the + // 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"); + printer_->Print("$alias$ = importlib.import_module('$name$')\n", "alias", + module_alias, "name", module_name); + } else { + int last_dot_pos = module_name.rfind('.'); + string import_statement; + if (last_dot_pos == string::npos) { + // NOTE(petya): this is not tested as it would require a protocol buffer + // outside of any package, and I don't think that is easily achievable. + import_statement = "import " + module_name; + } else { + import_statement = "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", module_alias); + } + CopyPublicDependenciesAliases(module_alias, file_->dependency(i)); } printer_->Print("\n"); diff --git a/src/google/protobuf/compiler/python/python_plugin_unittest.cc b/src/google/protobuf/compiler/python/python_plugin_unittest.cc index 23f2449c..34f857fd 100644 --- a/src/google/protobuf/compiler/python/python_plugin_unittest.cc +++ b/src/google/protobuf/compiler/python/python_plugin_unittest.cc @@ -46,6 +46,7 @@ #include <google/protobuf/testing/file.h> #include <google/protobuf/testing/file.h> +#include <google/protobuf/stubs/strutil.h> #include <google/protobuf/testing/googletest.h> #include <gtest/gtest.h> @@ -115,6 +116,53 @@ TEST(PythonPluginTest, PluginTest) { EXPECT_EQ(0, cli.Run(5, argv)); } +// This test verifies that the generated Python output uses regular imports (as +// opposed to importlib) in the usual case where the .proto file paths do not +// not contain any Python keywords. +TEST(PythonPluginTest, ImportTest) { + // Create files test1.proto and test2.proto with the former importing the + // latter. + GOOGLE_CHECK_OK(File::SetContents(TestTempDir() + "/test1.proto", + "syntax = \"proto3\";\n" + "package foo;\n" + "import \"test2.proto\";" + "message Message1 {\n" + " Message2 message_2 = 1;\n" + "}\n", + true)); + GOOGLE_CHECK_OK(File::SetContents(TestTempDir() + "/test2.proto", + "syntax = \"proto3\";\n" + "package foo;\n" + "message Message2 {}\n", + true)); + + google::protobuf::compiler::CommandLineInterface cli; + cli.SetInputsAreProtoPathRelative(true); + python::Generator python_generator; + cli.RegisterGenerator("--python_out", &python_generator, ""); + string proto_path = "-I" + TestTempDir(); + string python_out = "--python_out=" + TestTempDir(); + const char* argv[] = {"protoc", proto_path.c_str(), "-I.", python_out.c_str(), + "test1.proto"}; + ASSERT_EQ(0, cli.Run(5, argv)); + + // Loop over the lines of the generated code and verify that we find an + // ordinary Python import but do not find the string "importlib". + string output; + GOOGLE_CHECK_OK(File::GetContents(TestTempDir() + "/test1_pb2.py", &output, + true)); + std::vector<string> lines = Split(output, "\n"); + string expected_import = "import test2_pb2"; + bool found_expected_import = false; + for (int i = 0; i < lines.size(); ++i) { + if (lines[i].find(expected_import) != string::npos) { + found_expected_import = true; + } + EXPECT_EQ(string::npos, lines[i].find("importlib")); + } + EXPECT_TRUE(found_expected_import); +} + } // namespace } // namespace python } // namespace compiler |