String-based flags in python

For my python-based class I'd like to limit the input strings to certain strings, as they represent the file endings I want to allow to be loaded, i.e.

class DemoClass:
    def __init__(self, allowed_endings: Literal[".csv", ".txt", ".docx"] = ".docx") -> None:
        self.allowed_endings: Literal[".csv", ".txt", ".docx"] = allowed_endings

Now, I'd like to replace that implementation with a flag-based system, i.e.

class DemoClass:
    def __init__(self, allowed_endings: AllowedFileEndings = AllowedFileEndings.TXT) -> None:
        self.allowed_endings: AllowedFileEndings = allowed_endings

such that I can write a similar class as

class DemoClass:
    def __init__(self, allowed_endings: AllowedFileEndings = AllowedFileEndings.TXT | AllowedFileEndings.CSV) -> None:
        self.allowed_endings: AllowedFileEndings = allowed_endings

I am aware of StrEnum, where I can create a class which behaves like enums, but also can be used to be compared with strings. Unfortunately, it does not work with with bitwise OR/AND-operators, as demonstrated above. Moreover, it only works for python 3.11 and above, while I'm limited to 3.9.

Is there a way to implement such a StrFlag-class in Python 3.9?

Ideally, I can then still use string-based comprehensions, i.e.

".csv" is in AllowedFileEndings.TXT | AllowedFileEndings.CSV

if possible.

Solution

You can implement custom behavior for | using __or__ and __ror__ methods of a class. By implementing them on top of StrEnum you should get a class that fits your needs.

Since set already behave like that (| on sets is their union), they are perfect data structure for our usecase.

from enum import StrEnum

class Extensions(StrEnum):
    TXT = ".txt"
    CSV = ".csv"
    DOCX = ".docx"
    DOC = ".doc"

    def __or__(self, other):
        if isinstance(other, set):
            return {self} | other
        elif isinstance(other, Extensions):
            return {self, other}
        else:
            raise TypeError
    def __ror__(self, other):
        if isinstance(other, set):
            return other | {self}
        elif isinstance(other, Extensions):
            return {other, self}
        else:
            raise TypeError
        
print(Extensions.TXT | Extensions.CSV | Extensions.DOCX)

You have to remember though that while Extensions.TXT can be treated as a string, Extensions.TXT | Extensions.CSV becomes a set and must be treated accordingly.