Skip to content

Commit 35ac5f8

Browse files
serhiy-storchakaserhiy-storchaka
committed
Issue python#22955: attrgetter, itemgetter and methodcaller objects in the operator
module now support pickling. Added readable and evaluable repr for these objects. Based on patch by Josh Rosenberg.
1 parent 5418d0b commit 35ac5f8

4 files changed

Lines changed: 427 additions & 9 deletions

File tree

‎Lib/operator.py‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -231,65 +231,107 @@ class attrgetter:
231231
After h = attrgetter('name.first', 'name.last'), the call h(r) returns
232232
(r.name.first, r.name.last).
233233
"""
234+
__slots__ = ('_attrs', '_call')
235+
234236
def __init__(self, attr, *attrs):
235237
if not attrs:
236238
if not isinstance(attr, str):
237239
raise TypeError('attribute name must be a string')
240+
self._attrs = (attr,)
238241
names = attr.split('.')
239242
def func(obj):
240243
for name in names:
241244
obj = getattr(obj, name)
242245
return obj
243246
self._call = func
244247
else:
245-
getters = tuple(map(attrgetter, (attr,) + attrs))
248+
self._attrs = (attr,) + attrs
249+
getters = tuple(map(attrgetter, self._attrs))
246250
def func(obj):
247251
return tuple(getter(obj) for getter in getters)
248252
self._call = func
249253

250254
def __call__(self, obj):
251255
return self._call(obj)
252256

257+
def __repr__(self):
258+
return '%s.%s(%s)' % (self.__class__.__module__,
259+
self.__class__.__qualname__,
260+
', '.join(map(repr, self._attrs)))
261+
262+
def __reduce__(self):
263+
return self.__class__, self._attrs
264+
253265
class itemgetter:
254266
"""
255267
Return a callable object that fetches the given item(s) from its operand.
256268
After f = itemgetter(2), the call f(r) returns r[2].
257269
After g = itemgetter(2, 5, 3), the call g(r) returns (r[2], r[5], r[3])
258270
"""
271+
__slots__ = ('_items', '_call')
272+
259273
def __init__(self, item, *items):
260274
if not items:
275+
self._items = (item,)
261276
def func(obj):
262277
return obj[item]
263278
self._call = func
264279
else:
265-
items = (item,) + items
280+
self._items = items = (item,) + items
266281
def func(obj):
267282
return tuple(obj[i] for i in items)
268283
self._call = func
269284

270285
def __call__(self, obj):
271286
return self._call(obj)
272287

288+
def __repr__(self):
289+
return '%s.%s(%s)' % (self.__class__.__module__,
290+
self.__class__.__name__,
291+
', '.join(map(repr, self._items)))
292+
293+
def __reduce__(self):
294+
return self.__class__, self._items
295+
273296
class methodcaller:
274297
"""
275298
Return a callable object that calls the given method on its operand.
276299
After f = methodcaller('name'), the call f(r) returns r.name().
277300
After g = methodcaller('name', 'date', foo=1), the call g(r) returns
278301
r.name('date', foo=1).
279302
"""
303+
__slots__ = ('_name', '_args', '_kwargs')
280304

281305
def __init__(*args, **kwargs):
282306
if len(args) < 2:
283307
msg = "methodcaller needs at least one argument, the method name"
284308
raise TypeError(msg)
285309
self = args[0]
286310
self._name = args[1]
311+
if not isinstance(self._name, str):
312+
raise TypeError('method name must be a string')
287313
self._args = args[2:]
288314
self._kwargs = kwargs
289315

290316
def __call__(self, obj):
291317
return getattr(obj, self._name)(*self._args, **self._kwargs)
292318

319+
def __repr__(self):
320+
args = [repr(self._name)]
321+
args.extend(map(repr, self._args))
322+
args.extend('%s=%r' % (k, v) for k, v in self._kwargs.items())
323+
return '%s.%s(%s)' % (self.__class__.__module__,
324+
self.__class__.__name__,
325+
', '.join(args))
326+
327+
def __reduce__(self):
328+
if not self._kwargs:
329+
return self.__class__, (self._name,) + self._args
330+
else:
331+
from functools import partial
332+
return partial(self.__class__, self._name, **self._kwargs), self._args
333+
334+
293335
# In-place Operations *********************************************************#
294336

295337
def iadd(a, b):

‎Lib/test/test_operator.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import unittest
2+
import pickle
3+
import sys
24

35
from test import support
46

@@ -35,6 +37,9 @@ def __rmul__(self, other):
3537

3638

3739
class OperatorTestCase:
40+
def setUp(self):
41+
sys.modules['operator'] = self.module
42+
3843
def test_lt(self):
3944
operator = self.module
4045
self.assertRaises(TypeError, operator.lt)
@@ -396,6 +401,7 @@ def __getitem__(self, name):
396401
def test_methodcaller(self):
397402
operator = self.module
398403
self.assertRaises(TypeError, operator.methodcaller)
404+
self.assertRaises(TypeError, operator.methodcaller, 12)
399405
class A:
400406
def foo(self, *args, **kwds):
401407
return args[0] + args[1]
@@ -491,5 +497,108 @@ class PyOperatorTestCase(OperatorTestCase, unittest.TestCase):
491497
class COperatorTestCase(OperatorTestCase, unittest.TestCase):
492498
module = c_operator
493499

500+
501+
class OperatorPickleTestCase:
502+
def copy(self, obj, proto):
503+
with support.swap_item(sys.modules, 'operator', self.module):
504+
pickled = pickle.dumps(obj, proto)
505+
with support.swap_item(sys.modules, 'operator', self.module2):
506+
return pickle.loads(pickled)
507+
508+
def test_attrgetter(self):
509+
attrgetter = self.module.attrgetter
510+
attrgetter = self.module.attrgetter
511+
class A:
512+
pass
513+
a = A()
514+
a.x = 'X'
515+
a.y = 'Y'
516+
a.z = 'Z'
517+
a.t = A()
518+
a.t.u = A()
519+
a.t.u.v = 'V'
520+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
521+
with self.subTest(proto=proto):
522+
f = attrgetter('x')
523+
f2 = self.copy(f, proto)
524+
self.assertEqual(repr(f2), repr(f))
525+
self.assertEqual(f2(a), f(a))
526+
# multiple gets
527+
f = attrgetter('x', 'y', 'z')
528+
f2 = self.copy(f, proto)
529+
self.assertEqual(repr(f2), repr(f))
530+
self.assertEqual(f2(a), f(a))
531+
# recursive gets
532+
f = attrgetter('t.u.v')
533+
f2 = self.copy(f, proto)
534+
self.assertEqual(repr(f2), repr(f))
535+
self.assertEqual(f2(a), f(a))
536+
537+
def test_itemgetter(self):
538+
itemgetter = self.module.itemgetter
539+
a = 'ABCDE'
540+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
541+
with self.subTest(proto=proto):
542+
f = itemgetter(2)
543+
f2 = self.copy(f, proto)
544+
self.assertEqual(repr(f2), repr(f))
545+
self.assertEqual(f2(a), f(a))
546+
# multiple gets
547+
f = itemgetter(2, 0, 4)
548+
f2 = self.copy(f, proto)
549+
self.assertEqual(repr(f2), repr(f))
550+
self.assertEqual(f2(a), f(a))
551+
552+
def test_methodcaller(self):
553+
methodcaller = self.module.methodcaller
554+
class A:
555+
def foo(self, *args, **kwds):
556+
return args[0] + args[1]
557+
def bar(self, f=42):
558+
return f
559+
def baz(*args, **kwds):
560+
return kwds['name'], kwds['self']
561+
a = A()
562+
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
563+
with self.subTest(proto=proto):
564+
f = methodcaller('bar')
565+
f2 = self.copy(f, proto)
566+
self.assertEqual(repr(f2), repr(f))
567+
self.assertEqual(f2(a), f(a))
568+
# positional args
569+
f = methodcaller('foo', 1, 2)
570+
f2 = self.copy(f, proto)
571+
self.assertEqual(repr(f2), repr(f))
572+
self.assertEqual(f2(a), f(a))
573+
# keyword args
574+
f = methodcaller('bar', f=5)
575+
f2 = self.copy(f, proto)
576+
self.assertEqual(repr(f2), repr(f))
577+
self.assertEqual(f2(a), f(a))
578+
f = methodcaller('baz', self='eggs', name='spam')
579+
f2 = self.copy(f, proto)
580+
# Can't test repr consistently with multiple keyword args
581+
self.assertEqual(f2(a), f(a))
582+
583+
class PyPyOperatorPickleTestCase(OperatorPickleTestCase, unittest.TestCase):
584+
module = py_operator
585+
module2 = py_operator
586+
587+
@unittest.skipUnless(c_operator, 'requires _operator')
588+
class PyCOperatorPickleTestCase(OperatorPickleTestCase, unittest.TestCase):
589+
module = py_operator
590+
module2 = c_operator
591+
592+
@unittest.skipUnless(c_operator, 'requires _operator')
593+
class CPyOperatorPickleTestCase(OperatorPickleTestCase, unittest.TestCase):
594+
module = c_operator
595+
module2 = py_operator
596+
597+
@unittest.skipUnless(c_operator, 'requires _operator')
598+
class CCOperatorPickleTestCase(OperatorPickleTestCase, unittest.TestCase):
599+
module = c_operator
600+
module2 = c_operator
601+
602+
494603
if __name__ == "__main__":
495604
unittest.main()

‎Misc/NEWS‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ Core and Builtins
5252
Library
5353
-------
5454

55+
- Issue #22955: attrgetter, itemgetter and methodcaller objects in the operator
56+
module now support pickling. Added readable and evaluable repr for these
57+
objects. Based on patch by Josh Rosenberg.
58+
5559
- Issue #22107: tempfile.gettempdir() and tempfile.mkdtemp() now try again
5660
when a directory with the chosen name already exists on Windows as well as
5761
on Unix. tempfile.mkstemp() now fails early if parent directory is not

0 commit comments

Comments
 (0)