Do not check global module scope for Cython

1 files changed, 21 insertions, 15 deletions

diff --git a/python/google/protobuf/pyext/descriptor.cc b/python/google/protobuf/pyext/descriptor.cc
index 3526fc1c..12f3df41 100644
--- a/python/google/protobuf/pyext/descriptor.cc
+++ b/python/google/protobuf/pyext/descriptor.cc
@@ -101,20 +101,11 @@ bool _CalledFromGeneratedFile(int stacklevel) {
   if (frame == NULL) {
     return false;
   }
-  while (stacklevel-- > 0) {
-    frame = frame->f_back;
-    if (frame == NULL) {
-      return false;
-    }
-  }
-  if (frame->f_globals != frame->f_locals) {
-    // Not at global module scope
-    return false;
-  }
-
+  
   if (frame->f_code->co_filename == NULL) {
     return false;
   }
+  
   char* filename;
   Py_ssize_t filename_size;
   if (PyString_AsStringAndSize(frame->f_code->co_filename,
@@ -123,17 +114,32 @@ bool _CalledFromGeneratedFile(int stacklevel) {
     PyErr_Clear();
     return false;
   }
+  
+  if (filename_size < 3 or strcmp(&filename[filename_size - 3], ".py") != 0) {
+    // Cython is not using .py file and not at global module scope.
+    return true;
+  }
+  
   if (filename_size < 7) {
     // filename is too short.
     return false;
   }
-  // Cython is not using .py file. Only check filenames when end with .py
-  if (strcmp(&filename[filename_size - 3], ".py") == 0) {
-    if (strcmp(&filename[filename_size - 7], "_pb2.py") != 0) {
-      // Filename is not ending with _pb2.
+
+  if (strcmp(&filename[filename_size - 7], "_pb2.py") != 0) {
+    // Filename is not ending with _pb2.
+    return false;
+  }
+  
+  while (stacklevel-- > 0) {
+    frame = frame->f_back;
+    if (frame == NULL) {
       return false;
     }
   }
+  if (frame->f_globals != frame->f_locals) {
+    // Not at global module scope
+    return false;
+  }
 #endif
   return true;
 }