Commits

Marc Abramowitz committed d005382

Add six.map and six.filter, which always return a list, even in PY3.

Using six.map(f, seq) seems nicer than list(map(f, seq)). Ditto for filter.

Comments (0)

Files changed (2)

 else:
     wraps = functools.wraps
 
+def _wrap_builtin_func_return_with_list(func_name):
+    try:
+        # CPython
+        orig_func = __builtins__[func_name]
+    except TypeError:
+        # PyPy
+        orig_func = getattr(__builtins__, func_name)
+
+    if PY3:
+        @wraps(orig_func)
+        def _wrapper(*args, **kwargs):
+            return list(orig_func(*args, **kwargs))
+
+        return _wrapper
+    else:
+        return orig_func
+
+map = _wrap_builtin_func_return_with_list('map')
+filter = _wrap_builtin_func_return_with_list('filter')
+
 def with_metaclass(meta, *bases):
     """Create a base class with a metaclass."""
     # This requires a bit of explanation: the basic idea is to make a
     assert not hasattr(k, '__wrapped__')
 
 
+def test_six_map():
+    assert six.map(lambda x: x * 2, [1, 2, 3]) == [2, 4, 6]
+    assert six.map(lambda x, y: x + y, [1, 2, 3], [3, 2, 1]) == [4, 4, 4]
+    assert six.map.__doc__ == map.__doc__
+
+
+def test_six_filter():
+    assert six.filter(lambda x: x % 2 == 1, [1, 2, 3]) == [1, 3]
+    assert six.filter.__doc__ == filter.__doc__
+
+
 def test_add_metaclass():
     class Meta(type):
         pass