首页 > 解决方案 > Python类型:有没有办法为“元素”构造类型?

问题描述

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

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, 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)

标签: pythonpython-3.xpython-typing

解决方案


这对你有什么作用?

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,这样如果您不控制. 您应该能够很容易地添加参数。DiscreteBoxgym.sourcesshapeBox

这里的基本思想是我们Space用它可以包含的元素的类型进行参数化。我们使用NewType是为了使空间元素成为它们本质上的子类型(从中采样的元素Discreteint,并且具有int属性),而不会牺牲从 中采样的保证check强制执行。yspace


推荐阅读