Search code examples
pythonpython-typing

Python 3.6 type checking: numpy arrays and use defined classes


I'd think this has already been asked, but I can't find it

How can one check for non-builtin types in Python 3.6?

For example, let's say I want to have a function foo() that takes a numpy array as an arg, and returns an instance of my class Bah

def foo(a: np.array) -> Bah

Can I do something like this? How do I add those types to checking?


Solution

  • You can specify anything you want, as a type annotation if the type’s name is in scope (e.g., you’ve defined Bob in the current file, or done a from stuff importBob`).

    In general, the type checker doesn’t need any special knowledge of the type to know whether a value is that type. If it sees you store the result of this function in a variable whose type is Bob, or a supertype of Bob (which includes object and Any, and also includes unannotated variables), that’s legal; if it sees you store it in a variable whose type is int or some other unrelated type, it’s not. Similarly, if you pass the result on to some other function whose parameter is Bob or Any or unannotated, or you append it to a List[Bob], etc.

    But np.array is a different problem. That isn’t actually a type, it’s just a constructor function that usually returns a value of type np.ndarray, which is a type you don’t normally think about anywhere. So, a type checker can’t handle that without some kind of special information that array should be treated as a synonym for ndarray.

    Plus, many NumPy functions—and, implicitly, functions that you write yourself—actually take an “array-like”, which can be an ndarray, or usually a matrix, but also often any sequence. In which case you probably really want to either annotate then with something closer to accurate, like typing.Sequence—or maybe with a custom ArrayLike type.

    While we’re at it, you often want to specify the dtype—your function doesn’t want an array, it wants an array of floats, or it wants an array of <something> and wants to return a Bob whose values are that same <something>. So, you probably want a generic type like Sequence[float] or Sequence[T] to some typevar T.

    And you may even want to require a certain number of dimensions, or even a shape for those dimensions, or even a partial shape, or even that parameters x and y have to be broadcastable together or multiplyable. You can push this information into a generic type, but you’ll have to think it through pretty carefully.

    Anyway, for your own types, you rarely have to think that deep. Either Bob is a simple type that only needs the standard inheritance rules, so you don’t have to do anything, or it’s a generic collection type, where you just need to inherit/register it as a MutableSequence or a Mapping or whatever and it automatically gets the appropriate generic rules, or it’s a specific collection type, where you just inherit/register as a Set[int] and it automatically gets the approprIate rules.