When raising LDAPBytesWarning, walk the stack to determine stacklevel

tiran · encukou · commit a3723bc588ef · 2017-12-13T16:16:08.000+01:00
Closes: python-ldap#108 Signed-off-by: Christian Heimes <cheimes@redhat.com> Some simplification by: Petr Viktorin
diff --git a/Lib/ldap/ldapobject.py b/Lib/ldap/ldapobject.py
@@ -39,6 +39,9 @@
   text_type = str
 
 
+# See SimpleLDAPObject._bytesify_input
+_LDAP_WARN_SKIP_FRAME = True
+
 class LDAPBytesWarning(BytesWarning):
   """python-ldap bytes mode warning
   """
@@ -126,11 +129,27 @@ 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=6,
+            stacklevel=stacklevel+1,
           )
           return value.encode('utf-8')
     else:
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 SimpleLDAPObject._bytesify_input 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,64 @@ 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 _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)
+
+        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.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)
+        )
+
+    @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')
+
 
 class Test01_ReconnectLDAPObject(Test00_SimpleLDAPObject):
     """
