How can I inform mypy that a Set[re.Match] does not contain None?

Question

I have code:

matched_policies = {...} #some logic to derive a set
return {
        re.search(r"p\d{6}", matched_policy).group(0)
        for matched_policy in matched_policies
}

for which mypy throws error:

error: Item "None" of "Optional[Match[Any]]" has no attribute "group"

I know that there will be no None values in this set and I know from my previous question Can I inform mypy that an expression will not return an Optional? that its possible to inform mypy that a value will not be None by using an assert, but how do I do that for a set?

I tried changing my code to

assert None not in matched_policies
return {
    re.search(r_search, matched_policy).group(0)
    for matched_policy in matched_policies
}

but that made no difference, I still get the same error. I can get around it by using a loop:

prefixes = set()
for matched_policy in matched_policies:
    search_result = re.search(r_search, matched_policy)
    assert search_result is not None
    prefixes.add(search_result.group(0))
return prefixes

but that feels very inelegant so I don't want to do it, I would much rather use a set comprehension.

How can I assert that all elements in a set are not None in a way that satisfies mypy?

score 3 · Accepted Answer · answered Jul 22 '22 at 12:26

3

This issue is not with what's in your set, but that re.search returns Optional[Match[Any]].

If you have at least Python 3.8, you can use an assignment expression (aka walrus operator) inside your comprehension:

    matched_policies = {}
    return {
        match.group(0)
        for matched_policy in matched_policies
        if (match := re.search(r"p\d{6}", matched_policy)) is not None
    }

Mypy will find no issues there.

answered Jul 22 '22 at 12:26

ericbn

10,163
3
47
55

1

oh, this is great. This actually solves the performance issue I mentioned in my answer. – Berkay Berabi Jul 22 '22 at 12:29
nice. thank you very much (yes, we are using python 3.8) – jamiet Jul 22 '22 at 12:37

score 1 · Answer 2 · answered Jul 22 '22 at 12:14

I think you can accomplish what you want by using an if-clause in a list comprehension. Like below:

return {
        re.search(r"p\d{6}", matched_policy).group(0)
        for matched_policy in matched_policies if re.search(r_search, matched_policy) is not None
}

However, this comes at a performance cost because now you will compute the re.search two times. Once for the if condition and once for the actual operation.

If performance is not a concern, you can do that, but I would prefer the for loop you wrote.

How can I inform mypy that a Set[re.Match] does not contain None?

2 Answers2

Linked