Python类型：有没有一种方法可以为“元素”构造类型？

问题描述

我正在寻找一种创建类型的方法，该类型指示变量是某个其他集合的 元素。我知道集合类型：

from typing import Collection
Foo = Collection[Bar]

相反，我想做相反的事情，即

Bar = Element[Foo]

有没有办法做到这一点？

---

我想到的用例是能够实现以下功能：

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

---

以下是使用gym.spaces的详细示例：

from gym.spaces import Space,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)

解决方法

这如何为您工作？

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.

我将Space，Discrete和Box的类型提示写为type stubs，这样，如果您不控制{的来源，就可以添加它们{1}}。您应该可以轻松地将gym.sources参数添加到shape。

这里的基本思想是我们用其可以包含的元素类型对Box进行参数化。我们使用Space来使空间元素从根本上成为子类型（从NewType采样的元素是Discrete，并且具有int属性）而不会牺牲保证int强制执行，因为check是从y采样的。