Search code examples
pythonpycharmpython-typingmetaclass

Specify Typing of class with custom Metaclass


Following on a great system for using an enum-like replacement for Django choices (http://musings.tinbrain.net/blog/2017/may/15/alternative-enum-choices/) I have a project that uses a class with a custom metaclass that allows me to do list(MyChoices) (on the Class itself) to get a list of all the enum choices. The relevant part of the code looks something like this:

class MetaChoices(type):
    @classmethod
    def __prepare__(mcs, name, bases, **kwargs):
        return OrderedDict()

    def __new__(mcs, name, bases, attrs):
        _choices = OrderedDict()
        for attr_name, value in list(attrs.items()):
            ...do things...
        return type.__new__(mcs, name, bases, dict(attrs))

    def __iter__(cls):
        return iter(cls._choices.items())


class Choices(metaclass=MetaChoices):
    pass

class IceCreamFlavor(Choices):
    STRAWBERRY = ('strawberry', 'Fruity')
    CHOCOLATE = 'chocolate'

list(IceCreamFlavor)
# [('strawberry', 'Fruity'), ('chocolate', Chocolate')

The code has been working well for some time, but now I have typing turned on (in this case using PyCharm's type checker, but also looking for general solutions), and IceCreamFlavor is not marked as an iterable despite it being derived from a class whose metaclass defines the cls as having an __iter__ method. Does anyone know of a solution where I can show that the Choices class itself is itself an iterable?


Solution

  • I fixed the code to be correct for MyPy (checked easier by Pytype that adds annotation files *.pyi first).

    A typing problem was in the method __iter__(), that the attribute _choices seems undefined for a checker, because it was assigned not transparently, only by attrs['_choices'] = ....

    It can be annotated by adding one line:

    class MetaChoices(type):
        _choices = None  # type: dict   # written as comment for Python >= 3.5
        # _choices: dict                # this line can be uncommented if Python >= 3.6
    

    It is perfectly valid for Pytype and with its annotations it is checked valid also by MyPY of course.

    Maybe that typing problem in __iter__() could cause that the metaclass method was ignored in the checker.


    If the fix doesn't help, then the issue can be reported with the following simplified example:

    class MetaChoices(type):
        _choices = {0: 'a'}
    
        def __iter__(cls):
            return iter(cls._choices.items())
    
    
    class Choices(metaclass=MetaChoices):
        pass
    
    
    assert list(Choices) == [(0, 'a')]
    

    I reported another minor bug to the original article. That bug is not related to this problem.