Search code examples
pythonmypypython-typing

Constants as type hints for int bitflag

I have this code:

class Thing():
    FLAG_FOO = 1;
    FLAG_BAR = 2;
    FLAG_JOE = 4;

def call (self, flags: int) -> None:
    ...

The method Thing.call() should only accept one or a mix of the values that starts with FLAG_*, how can I put a type hint for that?

I would like, on VS Code or an IDE that supports type hints, when I write the function it display all the possible values available like so:

t = Thing();
t.call( ... );
        # popup menu that show options:
        # Thing.FLAG_FOO
        # Thing.FLAG_BAR
        # Thing.FLAG_JOE

I saw some modules can do that with literal strings but could this work with int bitflags?

Solution

  • Use enum.IntFlag:

    from enum import IntFlag

class Thing(IntFlag):
    FLAG_FOO = 1
    FLAG_BAR = 2
    FLAG_JOE = 4

    def call(self, flags: 'Thing') -> None:
        pass

    The type annotation specifies that you need to pass a Thing (note that for an Enum this means one of the class variables that you define as the enumeration; mypy understands how enums work and will error if you try to pass something else as the argument), and when you type Thing. you'll see the members in the IDE's autocomplete:

    Screenshot of VSCode