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"pd{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?

Asked By: jamiet

||

Source

Answer 1

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

return {
        re.search(r"pd{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.

Answered By: Berkay Berabi

Answer 2

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"pd{6}", matched_policy)) is not None
    }

Mypy will find no issues there.

Answered By: ericbn

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

Question:

Answers: