From 2d036025766bab417bb1d946f52918048e0f6589 Mon Sep 17 00:00:00 2001 From: Raymond Hettinger Date: Tue, 3 Feb 2009 02:28:00 +0000 Subject: [PATCH] Validate that __length_hint__ returns a usable result. --- Lib/test/test_iterlen.py | 10 ++++++++++ Objects/abstract.c | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/Lib/test/test_iterlen.py b/Lib/test/test_iterlen.py index 91873c2af4a..ab5466ecf99 100644 --- a/Lib/test/test_iterlen.py +++ b/Lib/test/test_iterlen.py @@ -208,6 +208,11 @@ class BadLengthHint(object): def __length_hint__(self): raise RuntimeError('hello') +class NoneLengthHint(object): + def __iter__(self): return iter(range(10)) + def __length_hint__(self): + return None + class TestLengthHintExceptions(unittest.TestCase): def test_issue1242657(self): @@ -225,6 +230,11 @@ class TestLengthHintExceptions(unittest.TestCase): self.assertRaises(RuntimeError, b.extend, BadLen()) self.assertRaises(RuntimeError, b.extend, BadLengthHint()) + def test_invalid_hint(self): + # Make sure an invalid result doesn't muck-up the works + self.assertEqual(list(NoneLengthHint()), list(range(10))) + + def test_main(): unittests = [ TestRepeat, diff --git a/Objects/abstract.c b/Objects/abstract.c index c10f4ca4f95..b47b1c9772f 100644 --- a/Objects/abstract.c +++ b/Objects/abstract.c @@ -123,7 +123,7 @@ _PyObject_LengthHint(PyObject *o, Py_ssize_t defaultvalue) PyErr_Clear(); return defaultvalue; } - rv = PyInt_AsLong(ro); + rv = rv = PyLong_Check(ro) ? PyLong_AsSsize_t(ro) : defaultvalue; Py_DECREF(ro); return rv; }