Python typing: is there a way to construct a type for "element of"?

Question

I'm looking for a way to create a type that indicates that a variable is an element of some other collection. I know of the Collection type:

from typing import Collection
Foo = Collection[Bar]

Instead, I'd like to do the inverse of that, i.e.

Bar = Element[Foo]

Is there a way to accomplish this?

---

The use-case I have in mind is to be able to something like:

import numpy as np
from gym.spaces import Space, Box, Discrete

Element = ...  # some type definition

def func(x: Element[Box], i: Element[Discrete]) -> Element[Box]:
    """ asserts are implied by the type annotations """
    assert isinstance(x, np.ndarray)
    assert isinstance(i, int)
    return x * i

---

Here's a slightly more detailed example using gym.spaces:

from gym.spaces import Space, Box, Discrete


box = Box(low=0, high=1, shape=(3,))
dsc = Discrete(5)

x = box.sample()  # example: x = array([0.917, 0.021, 0.740], dtype=float32)
i = dsc.sample()  # example: i = 3


def check(space: Space, y: Element[Space]) -> Element[Space]:
    if y not in space:
        raise ValueError("y not an element of space")
    return y


x = check(box, x)
i = check(dsc, i)

Answer 1

How does this work for you?

from abc import ABC, abstractmethod
from typing import Generic, TypeVar, NewType

T = TypeVar("T")
DiscreteT = NewType("DiscreteT", int)
BoxT = NewType("BoxT", float)

class Space(ABC, Generic[T]):
    @abstractmethod
    def sample(self) -> T: ...
    def __contains__(self, item: T) -> bool: ...

class Discrete(Space[DiscreteT]):
    def __init__(self, n: int) -> None: ...
    def sample(self) -> DiscreteT: ...

class Box(Space[BoxT]):
    def __init__(self, low: float, high: float) -> None: ...
    def sample(self) -> BoxT: ...

def check(space: Space[T], y: T) -> T:
    if y in space:
        raise ValueError("y not an element of space")

    return y

box = Box(low=0, high=1)
dsc = Discrete(5)

x = box.sample()
i = dsc.sample()

# Assumes that these lines are run separately for example's sake, such that assignment from one doesn't impact lines later.
x = check(box, x) # Passes mypy.
i = check(dsc, i) # Passes mypy.

x = check(box, i) # Fails mypy: error: Cannot infer type argument 1 of "check".
i = check(box, x) # Fails mypy: error: Cannot infer type argument 1 of "check".

other_dsc = Discrete(0)
i = check(other_dsc, i) # Passes mypy, even though `i` came from `dsc`. Don't know if it is possible for this to be caught at type-check time.

I have written the type hints for Space, Discrete and Box as type stubs, such that you can add them if you don't controll the source of gym.sources. You should be able to add the shape parameter to Box pretty easily.

The basic idea here is that we parametrize Space with the type of an element it can contain. We use NewType in order to make the space elements subtypes of what they fundamentally are (an element sampled from Discrete is an int, and has int properties) without sacrificing the guarantee check enforces in that y was sampled from space.