Search code examples
pythonpython-3.xpycharmmypypydantic

Dynamically updating type hints for all attributes of subclasses in pydantic


I am writing some library code where the purpose is to have a base data model that can be subclassed and used to implement data objects that correspond to the objects in a database. For this base model I am inheriting from pydantic.BaseModel.

There is a bunch of stuff going on but for this example essentially what I have is a base model class that looks something like this:

class Model(pydantic.BaseModel, metaclass=custom_complicated_metaclass):
    some_base_attribute: int
    some_other_base_attribute: str

I'll get back to what this metaclass does in a moment. This would then be subclassed by some user of this library like this:

class User(Model):
    age: int
    name: str
    birth_date: datetime.datetime

Now, the metaclass that I am using hooks in to getattr and allows the following syntax:

User.age > 18

Which then returns a custom Filter object that can be used to filter in a database, so essentially its using attribute access on the class directly (but not on its instances) as a way to create some syntactic sugar for user in filters and sorting.

Now, the issue comes when I would like to allow a number of attributes to sort the results of a database query by. I can do something like the following:

db = get_some_database(...)
db.query(..., order_by=[User.age, User.birth_date]

And this works fine, however I would like to be able to specify for each attribute in the order_by list if its ascending or descending order. The simplest syntax I could think of for this is to allow the use of - to invert the sort order like this:

db = get_some_database(...)
db.query(..., order_by=[User.age, -User.birth_date]

This works, I just implement __neg__ on my custom filter class and its all good.

Now finally, the issue I have is that because I defined User.birth_date to be a datetime object, it does not support the - operator which pycharm and mypy will complain about (and they will complain about it for any type that does not support -). They are kind of wrong since when accessing the attribute on the class like this instead of on an instance it actually will return an object that does support the - operator, but obviously they don't know this. If this would only be a problem inside my library code I wouldn't mind it so much, I could just ignore it or add a disable comment etc but since this false positive complaint will show up in end-user code I would really like to solve it.

So my actual question essentially is, can I in any way (that the type checkers would also understand) force all the attributes that are implemented on subclasses of my baseclass to have whatever type they are assigned but also union with my custom type, so that these complaints dont show up? Or is there another way I can solve this?


Solution

  • Plugin required

    Not without a custom plugin I'm afraid (see e.g. Pydantic).

    You explicitly annotate the birth_date name within the class' scope to be of the type datetime, so the type checkers are correct to say that it does not support __neg__ (instance or not is irrelevant).

    Your metaclass magic will likely not be understandable by type checkers, at least not today. This simple example demonstrates it:

    class MyMeta(type):
        def __getattr__(self, _: object) -> int:
            return 1
    
    
    class Model(metaclass=MyMeta):
        x: str
    
    
    print(Model.x)
    

    The output is 1 as expected, but adding reveal_type(Model.x) and running it through mypy reveals the type as builtins.str. This shows that mypy will choose the annotation in the class body regardless of what type MyMeta.__getattr__ returns.

    A union annotation is possible of course, but that will just cause problems on "the other end" because now mypy & Co. will never be sure, if the x attribute holds that special query-construct-type or the type of the actual value.


    Alternative A: Negation function

    If you are not willing to write your own type checker plugins (totally understandable), you could go with a less elegant workaround of a short-named negation function (in your query-sense). Something like this:

    from typing import Any
    
    
    class Special:
        val: int = 1
    
        def __repr__(self) -> str:
            return f"Special<{self.val}>"
    
    
    class MyMeta(type):
        def __getattr__(self, _: object) -> Special:
            return Special()
    
    
    class Model(metaclass=MyMeta):
        x: str
    
    
    def neg(obj: Any) -> Any:
        obj.val *= -1
        return obj
    
    
    print(Model.x, neg(Model.x))
    

    Output: Special<1> Special<-1>

    Not as nice as just prepending a -, but much less headache with type checkers.

    Note that Any here is actually important because otherwise you would just move the problem to the neg call site.


    Alternative B: Custom subtypes

    Alternatively you could of course subclass things like datetime and tell users to use your custom class for annotations instead.

    Then within the class body you only add a type stub for type checkers to satisfy their __neg__ complaints.

    Something like this:

    from __future__ import annotations
    from datetime import datetime
    from typing import TYPE_CHECKING
    
    
    class Special:
        val: int = 1
    
        def __repr__(self) -> str:
            return f"Special<{self.val}>"
    
        def __neg__(self) -> Special:
            obj = self.__class__()
            obj.val *= -1
            return obj
    
    
    class DateTime(datetime):
        if TYPE_CHECKING:
            def __neg__(self) -> Special:
                pass
    
    
    class MyMeta(type):
        def __getattr__(self, _: object) -> Special:
            return Special()
    
    
    class Model(metaclass=MyMeta):
        x: DateTime
    
    
    print(Model.x, -Model.x)
    

    Same output as before and type checkers don't complain about the - even though at runtime the DateTime and datetime classes are the same (though not identical).

    Again, not ideal because you would have to ask users to use a non-standard type here, but depending on the situation that may be a better option for you.

    Another drawback of this option compared to the negation function is that you would have to define a subtype for each type you want to allow in annotations. I don't know, if there will be many of those that don't support __neg__, but you would have to do this for all of them and users would have to use your types.