Pop specific item from Python set and get default if not exists

Question

Is there a simple method, maybe a one-liner, to pop an item from Python set, and get default value if it does not exist?

It could be done in four lines of code and two hash table lookups:

def pop_default(coll: set, v, default=None):
    if v not in coll:
        return default
    coll.remove(v)
    return v

Can this function be re-written and

1. Be more compact, preferably a one-liner,
2. Be at least same efficient?

Function usage example:

s = {'a', 'b', 'c'}
assert pop_default(s, 'a') == 'a'
assert s == {'b', 'c'}

assert pop_default(s, 'x') == None
assert s == {'b', 'c'}

Solution

Two separate improvements were proposed by @jonrsharpe and @TimurShtatland.

One-liner

def pop_default(s: set, v, default=None):
    return s.remove(v) or v if v in s else default

A clear winner. Same efficiency, one line of code.

Alternative for some rare cases: exception-based

def pop_default(coll: set, v, default=None):
    try:
        coll.remove(v)
        return v
    except KeyError:
        return default

This solution is ~5 times slower in @nocomment 's test but is faster if exception is not raised.

Answer 1

This requires only a single lookup into the set, so in rare corner cases it is somewhat more efficient. Unfortunately, more generally (which includes the OP case, as well as many realistic cases), it is less efficient. Also, it is longer by one line. Note that I renamed the method to avoid clashing with the existing pop method.

def remove_with_default(coll: set, v, default=None):
    try:
        coll.remove(v)
        return v
    except KeyError:
        return default

s = {'a', 'b', 'c'}
assert remove_with_default(s, 'a') == 'a'
assert s == {'b', 'c'}

assert remove_with_default(s, 'x') == None
assert s == {'b', 'c'}

Benchmark

(Thanks to no comment for the benchmarking code and to the OP for the suggestion of the corner case)

from timeit import repeat
from random import randrange
import numpy as np

def pop(coll: set, v, default=None):
    if v not in coll:
        return default
    coll.remove(v)
    return v

def remove_with_default(coll: set, v, default=None):
    try:
        coll.remove(v)
        return v
    except KeyError:
        return default

set_size = 10
num_reps = 10**5

print("OP case:")
for f in [pop, remove_with_default] * 3:
    s = {'a', 'b', 'c'}
    t = 1e9 / num_reps * np.mean(repeat("f(s, 'a')", number=num_reps, globals=globals()))
    print(f'{t:0.1e} ns ', f.__name__)

print("more realistic case:")
for f in [pop, remove_with_default] * 3:
    t = 1e9 / num_reps * np.mean(repeat("s = set(range(set_size)); [f(s, randrange(set_size)) for i in range(set_size)]", number=num_reps, globals=globals()))
    print(f'{t:0.1e} ns ', f.__name__)

print("corner case:")
for f in [pop, remove_with_default] * 3:
    t = 1e9 / num_reps * np.mean(repeat("s = set(range(set_size)); [f(s, i) for i in range(set_size)]", number=num_reps, globals=globals()))
    print(f'{t:0.1e} ns ', f.__name__)

Results:

OP case:
6.4e+01 ns  pop
1.7e+02 ns  remove_with_default
3.3e+01 ns  pop
1.6e+02 ns  remove_with_default
3.3e+01 ns  pop
1.6e+02 ns  remove_with_default
more realistic case:
2.8e+03 ns  pop
3.4e+03 ns  remove_with_default
2.8e+03 ns  pop
3.3e+03 ns  remove_with_default
2.8e+03 ns  pop
3.4e+03 ns  remove_with_default
corner case:
8.2e+02 ns  pop
7.4e+02 ns  remove_with_default
8.2e+02 ns  pop
7.1e+02 ns  remove_with_default
8.5e+02 ns  pop
7.3e+02 ns  remove_with_default