pythonenumspython-typingpyright

How to strongly type the "value" attribute to be str or a custom type?


When using the Pylance (ms-python.vscode-pylance) VS Code extension in strict type checking mode, I get a type error on my custom Enum value for the following code:

def println_ctrl_sequence(message: str, ctrlSequence: Union[ANSICtrlSequence, str]):
    """
    This function is use with  terminals to print the message
    with colors specified by a, ANSI control sequence that 
    can be either a str or a console.ANSICtrlSequence object.
    """
    if type(ctrlSequence) == ANSICtrlSequence:
        ctrlSequenceStr: str = ctrlSequence.value
    else:
        ctrlSequenceStr = ctrlSequence
    
    print("%s%s%s" % (
        ctrlSequenceStr,
        message,
        ANSICtrlSequence.RESET.value
    ))

The type error is detected on the ctrlSequenceStr: str = ctrlSequence.value line since ctrlSequence.valueis detected as being of type Any | Unknown. So my objective is to strongly type the value attribute of my extended Enum:

# python enum : https://docs.python.org/3/library/enum.html 
from enum import Enum

class ANSICtrlSequence(Enum):

    # basic control sequences
    RESET = "\033[m" 

    # full control sequences
    PASSED = "\033[1;4;38;5;76m" 
    FAILED = "\033[1;5;38;5;197m" 

I have tried things like for instance doing ANSICtrlSequence(str, Enum) as specified here in "String-based enum in Python" Q&A without success.

I have read the class enum.pyi and I can understand why the type of value is what it is:

class Enum(metaclass=EnumMeta):
    name: str
    value: Any
    ...

I can't find a way to type my value attribute to be a str anywhere in the documentation or on StackOverflow. So is it possible? Is there a way to override the type of an inherited attribute? Or do I need to extend the Enum class with for instance an equivalent of the IntEnum that could be StrEnum for instance? Maybe I need to write my own strongly typed Enum class? Is there anything I missed?


Solution

  • It seems the problem comes not entirely from Enum.value but from including str as a possible type of ctrlSequence. Pylance seems to check if all the types included in the Union has a .value attribute, and for str, of course, it doesn't have .value so Pylance doesn't know what type to expect (that's the "Unknown").

    We can reproduce a similar "Unknown" error without using an Enum:

    x = 5
    print(x.value)
    

    similar error without enum

    In your case, following the string-based enum solution and inheriting from str when defining your Enum is still necessary, as that indicates to type-checkers (here, Pylance) that your Enum .value's are of str-type.

    So, you definitely need this to strongly type your Enum:

    class ANSICtrlSequence(str, Enum):
        RESET = "\033[m" 
        PASSED = "\033[1;4;38;5;76m" 
        FAILED = "\033[1;5;38;5;197m" 
    

    But then, it still shows up as

    Type of "value" is partially unknown
    ^ Type of "value" is "Any | Unknown*"

    because in Union[ANSICtrlSequence, str], the type of .value for ANSICtrlSequence is Any and the type of .value for str is Unknown. This issue with str is evident when you reverse the order of the Union to Union[str, ANSICtrlSequence], which then becomes

     Type of "value" is "Unknown | Any"

    ...indicating the "Unknown" is tied to the str. Basically, my point is you should not be focusing on typing the .value attribute of the Enum, because the problem is with including str. The error actually disappears if you remove the Union and just use ANSICtrlSequence:

    class ANSICtrlSequence(str, Enum):
        RESET = "\033[m" 
        PASSED = "\033[1;4;38;5;76m" 
        FAILED = "\033[1;5;38;5;197m" 
    
    def println_ctrl_sequence(message: str, ctrlSequence: ANSICtrlSequence):
        # Pylance does not complain here
        ctrlSequenceStr: str = ctrlSequence.value
    

    ...which indicates there isn't any problem with your Enum in the first place.

    But I understand why there's a Union in the code. Unfortunately, Pylance doesn't. It doesn't understand that by the time the code gets to ctrlSequence.value the code already checked that ctrlSequence is an Enum.

    Interestingly, what does work is to change how to check for the type. Instead of type(obj), use isinstance(obj, classinfo):

    class ANSICtrlSequence(str, Enum):
        RESET = "\033[m" 
        PASSED = "\033[1;4;38;5;76m" 
        FAILED = "\033[1;5;38;5;197m" 
    
    def println_ctrl_sequence(message: str, ctrlSequence: Union[ANSICtrlSequence, str]):
        if isinstance(ctrlSequence, ANSICtrlSequence):
            ctrlSequenceStr: str = ctrlSequence.value
        else:
            ctrlSequenceStr = ctrlSequence
    

    no errors with isinstance

    ...which satisfies Pylance and fixes the errors :)

    I hope this isn't a case of "it works on my environment", but I almost always use isinstance instead of type when checking for an object's type, and I don't get any Pylance errors with Enums or Unions with Enums in my codes. I don't know how Pylance works but here's a related Q&A on the topic: What are the differences between type() and isinstance()?

    If you really need to use type(ctrlSequence), then you can use typing.cast, which:

    Cast a value to a type.

    This returns the value unchanged. To the type checker this signals that the return value has the designated type, but at runtime we intentionally don’t check anything (we want this to be as fast as possible).

    from typing import Union, cast
    
    class ANSICtrlSequence(str, Enum):
        RESET = "\033[m" 
        PASSED = "\033[1;4;38;5;76m" 
        FAILED = "\033[1;5;38;5;197m" 
    
    def println_ctrl_sequence(message: str, ctrlSequence: Union[ANSICtrlSequence, str]):
        if type(ctrlSequence) == ANSICtrlSequence:
            ctrlSequence = cast(ANSICtrlSequence, ctrlSequence)
            ctrlSequenceStr: str = ctrlSequence.value
        else:
            ctrlSequenceStr = ctrlSequence
    

    no errors by using cast

    ...which again satisfies Pylance and fixes the errors :) The cast enforces the type-checker (here, Pylance) that ctrlSequence is your Enum type, and the .value is indeed a string.