Use stack walking for LDAPBytesWarning also in initialize()

encukou · encukou · commit db98910a3a3c · 2017-12-13T16:16:08.000+01:00
diff --git a/Lib/ldap/functions.py b/Lib/ldap/functions.py
@@ -27,6 +27,9 @@
   # Tracing is only supported in debugging mode
   import traceback
 
+# See _raise_byteswarning in ldapobject.py
+_LDAP_WARN_SKIP_FRAME = True
+
 
 def _ldap_function_call(lock,func,*args,**kwargs):
   """
diff --git a/Lib/ldap/ldapobject.py b/Lib/ldap/ldapobject.py
@@ -46,6 +46,26 @@ class LDAPBytesWarning(BytesWarning):
   """python-ldap bytes mode warning
   """
 
+def _raise_byteswarning(message):
+    """Raise LDAPBytesWarning
+    """
+
+    # Call stacks that raise the warning tend to be complicated, so
+    # getting a useful stacklevel is tricky.
+    # We walk stack frames, ignoring functions in uninteresting files,
+    # based on the _LDAP_WARN_SKIP_FRAME marker in globals().
+    stacklevel = 2
+    try:
+        getframe = sys._getframe
+    except AttributeError:
+        pass
+    else:
+        frame = sys._getframe(stacklevel)
+        while frame and frame.f_globals.get('_LDAP_WARN_SKIP_FRAME'):
+            stacklevel += 1
+            frame = frame.f_back
+    warnings.warn(message, LDAPBytesWarning, stacklevel=stacklevel+1)
+
 
 class NO_UNIQUE_ENTRY(ldap.NO_SUCH_OBJECT):
   """
@@ -90,13 +110,10 @@ def __init__(
     # By default, raise a TypeError when receiving invalid args
     self.bytes_mode_hardfail = True
     if bytes_mode is None and PY2:
-      warnings.warn(
+      _raise_byteswarning(
         "Under Python 2, python-ldap uses bytes by default. "
         "This will be removed in Python 3 (no bytes for DN/RDN/field names). "
-        "Please call initialize(..., bytes_mode=False) explicitly.",
-        LDAPBytesWarning,
-        stacklevel=2,
-      )
+        "Please call initialize(..., bytes_mode=False) explicitly.")
       bytes_mode = True
       # Disable hard failure when running in backwards compatibility mode.
       self.bytes_mode_hardfail = False
@@ -129,28 +146,10 @@ def _bytesify_input(self, value):
         if self.bytes_mode_hardfail:
           raise TypeError("All provided fields *must* be bytes when bytes mode is on; got %r" % (value,))
         else:
-          # Raise LDAPBytesWarning.
-          # Call stacks with _bytesify_input tend to be complicated, so
-          # getting a useful stacklevel is tricky.
-          # We walk stack frames, ignoring all functions in this file
-          # and in the _ldap extension, based on a marker in globals().
-          stacklevel = 0
-          try:
-            getframe = sys._getframe
-          except AttributeError:
-            pass
-          else:
-            frame = sys._getframe(stacklevel)
-            # walk up the stacks until we leave the file
-            while frame and frame.f_globals.get('_LDAP_WARN_SKIP_FRAME'):
-                stacklevel += 1
-                frame = frame.f_back
-          warnings.warn(
-            "Received non-bytes value %r with default (disabled) bytes mode; please choose an explicit "
-            "option for bytes_mode on your LDAP connection" % (value,),
-            LDAPBytesWarning,
-            stacklevel=stacklevel+1,
-          )
+          _raise_byteswarning(
+            "Received non-bytes value %r with default (disabled) bytes mode; "
+            "please choose an explicit "
+            "option for bytes_mode on your LDAP connection" % (value,))
           return value.encode('utf-8')
     else:
       if not isinstance(value, text_type):
diff --git a/Modules/ldapmodule.c b/Modules/ldapmodule.c
@@ -73,7 +73,7 @@ PyObject* init_ldap_module(void)
 	LDAPinit_control(d);
 
     /* Marker for LDAPBytesWarning stack walking
-     * see SimpleLDAPObject._bytesify_input in ldapobject.py
+     * See _raise_byteswarning in ldapobject.py
      */
     if (PyModule_AddIntConstant(m, "_LDAP_WARN_SKIP_FRAME", 1) != 0) {
         return NULL;
diff --git a/Tests/t_ldapobject.py b/Tests/t_ldapobject.py
@@ -361,6 +361,25 @@ def catch_byteswarnings(self, *args, **kwargs):
             warnings.simplefilter('always', ldap.LDAPBytesWarning)
             yield conn, w
 
+    def _check_byteswarning(self, warning, expected_message):
+        self.assertIs(warning.category, ldap.LDAPBytesWarning)
+        self.assertIn(expected_message, text_type(warning.message))
+
+        def _normalize(filename):
+            # Python 2 likes to report the ".pyc" file in warnings,
+            # tracebacks or __file__.
+            # Use the corresponding ".py" in that case.
+            if filename.endswith('.pyc'):
+                return filename[:-1]
+            return filename
+
+        # Assert warning points to a line marked CORRECT LINE in this file
+        self.assertEquals(_normalize(warning.filename), _normalize(__file__))
+        self.assertIn(
+            'CORRECT LINE',
+            linecache.getline(warning.filename, warning.lineno)
+        )
+
     def _test_byteswarning_level_search(self, methodname):
         with self.catch_byteswarnings(explicit=False) as (conn, w):
             method = getattr(conn, methodname)
@@ -374,43 +393,31 @@ def _test_byteswarning_level_search(self, methodname):
 
         self.assertEqual(len(w), 2, w)
 
-        def _normalize(filename):
-            # Python 2 likes to report the ".pyc" file in warnings,
-            # tracebacks or __file__.
-            # Use the corresponding ".py" in that case.
-            if filename.endswith('.pyc'):
-                return filename[:-1]
-            return filename
-
-        self.assertIs(w[0].category, ldap.LDAPBytesWarning)
-        self.assertIn(
-            u"Received non-bytes value u'(cn=Foo*)'",
-            text_type(w[0].message)
-        )
-        self.assertEqual(_normalize(w[1].filename), _normalize(__file__))
-        self.assertEqual(_normalize(w[0].filename), _normalize(__file__))
-        self.assertIn(
-            'CORRECT LINE',
-            linecache.getline(w[0].filename, w[0].lineno)
-        )
+        self._check_byteswarning(
+            w[0], u"Received non-bytes value u'(cn=Foo*)'")
 
-        self.assertIs(w[1].category, ldap.LDAPBytesWarning)
-        self.assertIn(
-            u"Received non-bytes value u'*'",
-            text_type(w[1].message)
-        )
-        self.assertIn(_normalize(w[1].filename), _normalize(__file__))
-        self.assertIn(
-            'CORRECT LINE',
-            linecache.getline(w[1].filename, w[1].lineno)
-        )
+        self._check_byteswarning(
+            w[1], u"Received non-bytes value u'*'")
 
     @unittest.skipUnless(PY2, "no bytes_mode under Py3")
     def test_byteswarning_level_search(self):
         self._test_byteswarning_level_search('search_s')
         self._test_byteswarning_level_search('search_st')
         self._test_byteswarning_level_search('search_ext_s')
 
+    @unittest.skipUnless(PY2, "no bytes_mode under Py3")
+    def test_byteswarning_initialize(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.resetwarnings()
+            warnings.simplefilter('always', ldap.LDAPBytesWarning)
+            bytes_uri = self.server.ldap_uri.decode('utf-8')
+            self.ldap_object_class(bytes_uri)  # CORRECT LINE
+
+        self.assertEqual(len(w), 1, w)
+
+        self._check_byteswarning(
+            w[0], u"Under Python 2, python-ldap uses bytes by default.")
+
 
 class Test01_ReconnectLDAPObject(Test00_SimpleLDAPObject):
     """
