Shared via mypy Playground

main.py

from enum import StrEnum, auto
from dataclasses import dataclass

class E(StrEnum):
    A = auto()
    B = auto()
    
@dataclass()
class D:
    e: E
    
    def set_e(self, e: E) -> None:
        self.e = e
    
def func(d: D, e: E) -> None:
    reveal_type(d.e)
    if d.e == E.A:
        return
    d.set_e(e)
    reveal_type(d.e)
    if d.e == E.A:
        return
    reveal_type(d.e)