Merge pull request python-ldap#128 – When raising LDAPBytesWarning, walk the stack to determine stacklevel

encukou · web-flow · commit 4bc1f58aa405 · 2017-12-15T11:59:46.000+01:00
python-ldap#128
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
@@ -39,10 +39,33 @@
   text_type = str
 
 
+# See SimpleLDAPObject._bytesify_input
+_LDAP_WARN_SKIP_FRAME = True
+
 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):
   """
@@ -87,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
@@ -126,12 +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:
-          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=6,
-          )
+          _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
@@ -72,6 +72,13 @@ PyObject* init_ldap_module(void)
 	LDAPinit_functions(d);
 	LDAPinit_control(d);
 
+    /* Marker for LDAPBytesWarning stack walking
+     * See _raise_byteswarning in ldapobject.py
+     */
+    if (PyModule_AddIntConstant(m, "_LDAP_WARN_SKIP_FRAME", 1) != 0) {
+        return NULL;
+    }
+
 	/* Check for errors */
 	if (PyErr_Occurred())
 		Py_FatalError("can't initialize module _ldap");
diff --git a/Tests/t_ldapobject.py b/Tests/t_ldapobject.py
@@ -16,8 +16,11 @@
     PY2 = False
     text_type = str
 
+import contextlib
+import linecache
 import os
 import unittest
+import warnings
 import pickle
 import warnings
 from slapdtest import SlapdTestCase, requires_sasl
@@ -329,7 +332,7 @@ def test_ldapbyteswarning(self):
         self.assertIsInstance(self.server.suffix, text_type)
         with warnings.catch_warnings(record=True) as w:
             warnings.resetwarnings()
-            warnings.simplefilter('default')
+            warnings.simplefilter('always', ldap.LDAPBytesWarning)
             conn = self._get_bytes_ldapobject(explicit=False)
             result = conn.search_s(
                 self.server.suffix,
@@ -350,6 +353,71 @@ def test_ldapbyteswarning(self):
             "LDAP connection" % self.server.suffix
         )
 
+    @contextlib.contextmanager
+    def catch_byteswarnings(self, *args, **kwargs):
+        with warnings.catch_warnings(record=True) as w:
+            conn = self._get_bytes_ldapobject(*args, **kwargs)
+            warnings.resetwarnings()
+            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)
+            result = method(
+                self.server.suffix.encode('utf-8'),
+                ldap.SCOPE_SUBTREE,
+                '(cn=Foo*)',
+                attrlist=['*'],  # CORRECT LINE
+            )
+            self.assertEqual(len(result), 4)
+
+        self.assertEqual(len(w), 2, w)
+
+        self._check_byteswarning(
+            w[0], u"Received non-bytes value u'(cn=Foo*)'")
+
+        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):
     """
