Adding attributes to Python's typing types

Question

I am trying to create types with certain constraints.

Because I want these constraints to be arbitrarily complex, I've decided that I do not need them to be type-checked, but I do want them to travel with the type definitions.

As an example, say I want to create a homogeneous container (i.e. typing.List) constrained to a size. For this particular example, I know I should use Tuple[int, int] (as per this question) but that's not flexible enough for other use cases.

My desired functionality:

from typing import List

class MyList(List):
    def __init__(self, *, num_elements: int) -> None:
        self.num_elements = num_elements

    def validate(self, input: List) -> None:
        if len(to_validate) > self.num_elements:
            raise ValueError


class MyClass:
    myvar: MyList(num_elements=2)[int]  # should look like List[int] to a type checker

    def __init__(self, *, myvar: MyList(num_elements=2)[int]):  # I guess I need to define twice?
        self.myvar = myvar
        self.validate_all()
    
    def validate_all(self):
        for var in self.__annotations__:
            if hasattr(self.__annotations__[var], "validate"):
                self.__annotations__[var].validate(getattr(self, var))


MyClass(myvar=(1, 2))  # pass

MyClass(myvar=(1, 2, 3))  # fail

As noted above, the annotation for myvar would have to look like List[int] to a type checker like mypy so that at least that part can be handled by existing frameworks.

I know I probably need to do something with typing.Generic or typing.TypeVar but I tried and honestly don't understand how they would be applicable to this situation.

Answer 1

I think I found a solution using pydantic (credit to Tomasz Wojcik for suggesting it in their answer).

This is a trivial example for a fixed size list. I'm using this example because it is what was in my original question, however I will note that pydantic has this constrained type built in (see their docs). Still, this general approach does allow arbitrary constrained types that pydantic will validate.

I do not know how this will interact with other type-checking systems (thinking of mypy), but since pydantic is compatible with mypy I think this should too.

Most of this code/solution was borrowed from the implementation of pydantic's existing constrained types (here). I would suggest that anyone trying to apply this method take a look there for a starting point.

from typing import Optional, List, TypeVar, Type, Dict, Any, Generator, Callable
from types import new_class

from pydantic import BaseModel
from pydantic.error_wrappers import ValidationError
from pydantic.fields import ModelField
from pydantic.utils import update_not_none
from pydantic.validators import list_validator

T = TypeVar('T')

# This types superclass should be List[T], but cython chokes on that...
class ConstrainedListClass(list):  # type: ignore
    # Needed for pydantic to detect that this is a list
    __origin__ = list
    __args__: List[Type[T]]  # type: ignore
    item_type: Type[T]  # type: ignore

    length: Optional[int] = None
    
    @classmethod
    def __get_validators__(cls) -> Generator[Callable, None, None]:
        yield cls.validator

    @classmethod
    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
        update_not_none(field_schema, length=cls.length)

    @classmethod
    def validator(cls, v: 'Optional[List[T]]', field: 'ModelField') -> 'Optional[List[T]]':
        if v is None and not field.required:
            return None

        v = list_validator(v)

        if cls.length != len(v):
            raise ValueError

        return v

def ConstrainedList(item_type: Type[T], *, length: int = None) -> Type[List[T]]:
    """Factory function for ConstrainedListClass

    Returns
    -------
    Type[List[T]]
        [description]
    """
    cls = ConstrainedListClass
    # __args__ is needed to conform to typing generics api
    namespace = {'length': length, 'item_type': item_type, '__args__': [item_type]}
    # We use new_class to be able to deal with Generic types
    return new_class(name=cls.__name__ + 'Value', bases=(cls, ), kwds={}, exec_body=lambda ns: ns.update(namespace))

ConstrainedList(int, length=2)

class UserClass(BaseModel):
    myvar: ConstrainedList(int, length=2)


UserClass(myvar=[1, 2])   # pass

for myvar in [
    [1],  # wrong number of items
    ["str1", "str2"],   # wrong type
]:
    try:
        UserClass(myvar=myvar)  # fail
    except ValidationError:
        continue
    raise RuntimeError

I will also note that it would be nice if the factory function could be put into the class' __new__ so that the syntax could stay the same but avoid having the extra function. I couldn't get this or anything similar to work.

Answer 2

This doesn't use typing, but you could solve this using the descriptor pattern. This would be something like this:

from collections import Iterable


class MyList:
    def __init__(self, name):
        self.name = name

    def __get__(self, instance, cls):
        if instance is None:
            return self
        else:
            return instance.__dict__[self.name]

    def __set__(self, instance, value):
        if not isinstance(value, Iterable):
            raise TypeError("This isn't iterable!")
        elif len(value) != 2 or not all(isinstance(val, int) for val in value):
            raise TypeError("List must be length 2 and all integers")

        instance.__dict__[self.name] = value

    def __delete__(self, instance):
        del instance.__dict__[self.name]


class MyClass:
    variable = MyList("variable")
    def __init__(self, var):
        self.variable = var


MyClass(myvar=(1, 2))  # pass
MyClass(myvar=[1, 2])  # pass
MyClass(myvar=(1, 2, 3))  # fail
MyClass(myvar="1,2,3")  # fail

Answer 3

To my knowledge typing module is quite primitive because we needed a basic type checking with some extra flavors and that's all it does - allows to valiadate the type of an input. What you want is a logical validation that is't really defined by a type. List with 3 or 2 elements is still a list.

With pydantic you can do

from typing import List

from pydantic import validator, BaseModel

num_elements = 2


class MyClass(BaseModel):
    myvar: List[int]

    @validator('myvar')
    def check_myvar_length(cls, v):
        if len(v) > num_elements:
            raise ValueError("Myvar too long!")
        return v

MyClass(myvar=(1, 2))  # pass
MyClass(myvar=(1, 2, 3))  # fail

or with dataclasses

from dataclasses import dataclass
from typing import List

num_elements = 2


@dataclass
class MyClass:
    myvar: List[int]

    def __post_init__(self):
        if len(self.myvar) > num_elements:
            raise ValueError("Myvar too long!")


MyClass(myvar=(1, 2))  # pass
MyClass(myvar=(1, 2, 3))  # fail

I know what you want to accomplish but I don't think it's possible. You can always create a regular class with validate method and run it in __init__ but I don't think that's what you want nor it is readable.