From 0442ed62cbea989742933d6c90f4f0a70078959a Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Wed, 20 Mar 2024 14:28:50 -0700
Subject: [PATCH] Enable Protobuf Python version check.

PiperOrigin-RevId: 617626372
---
 python/google/protobuf/runtime_version.py     |  2 +-
 .../protobuf/compiler/python/generator.cc     | 27 +++++++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/python/google/protobuf/runtime_version.py b/python/google/protobuf/runtime_version.py
index 2f2464b667af..519b1a308810 100644
--- a/python/google/protobuf/runtime_version.py
+++ b/python/google/protobuf/runtime_version.py
@@ -89,7 +89,7 @@ def ValidateProtobufRuntimeVersion(
         f' cannot be older than the linked gencode version. {error_prompt}'
     )
 
-  if gen_suffix is not SUFFIX:
+  if gen_suffix != SUFFIX:
     raise VersionError(
         'Detected mismatched Protobuf Gencode/Runtime version suffixes when'
         f' loading {location}: gencode {gen_version} runtime {version}.'
diff --git a/src/google/protobuf/compiler/python/generator.cc b/src/google/protobuf/compiler/python/generator.cc
index fbfe5fd56249..03f21ba00ddd 100644
--- a/src/google/protobuf/compiler/python/generator.cc
+++ b/src/google/protobuf/compiler/python/generator.cc
@@ -367,12 +367,35 @@ void Generator::PrintTopBoilerplate() const {
     // instead uses aliases assigned when importing modules.
     printer_->Print("import google3\n");
   }
+  bool runtime_version_disabled = false;
   printer_->Print(
       "from google.protobuf import descriptor as _descriptor\n"
       "from google.protobuf import descriptor_pool as _descriptor_pool\n"
+      "$runtime_version_import$"
       "from google.protobuf import symbol_database as _symbol_database\n"
-      "from google.protobuf.internal import builder as _builder\n");
-
+      "from google.protobuf.internal import builder as _builder\n",
+      "runtime_version_import",
+      runtime_version_disabled ? ""
+                               : "from google.protobuf import runtime_version "
+                                 "as _runtime_version\n");
+  if (!runtime_version_disabled) {
+    const auto& version = GetProtobufPythonVersion(opensource_runtime_);
+    printer_->Print(
+        "_runtime_version.ValidateProtobufRuntimeVersion(\n"
+        "    $domain$,\n"
+        "    $major$,\n"
+        "    $minor$,\n"
+        "    $patch$,\n"
+        "    '$suffix$',\n"
+        "    '$location$'\n"
+        ")\n",
+        "domain",
+        opensource_runtime_ ? "_runtime_version.Domain.PUBLIC"
+                            : "_runtime_version.Domain.GOOGLE_INTERNAL",
+        "major", absl::StrCat(version.major()), "minor",
+        absl::StrCat(version.minor()), "patch", absl::StrCat(version.patch()),
+        "suffix", version.suffix(), "location", file_->name());
+  }
   printer_->Print("# @@protoc_insertion_point(imports)\n\n");
   printer_->Print("_sym_db = _symbol_database.Default()\n");
   printer_->Print("\n\n");