How to define a type alias (TypeAlias) that is evaluated lazily via ForwardRef in python?

566 Views Asked by At

My goal is to prevent the import of expensive modules at runtime when using a type alias.

Without an alias one can hide expensive modules behind typing.TYPE_CHECKING and make the import local to the respective functions:

# utils.py -- module that must be very lightweight to import
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    import tensorflow as tf  # lot's of unwanted side-effects

def rarely_used_function(t: "tf.Tensor"):  # <-- note the lazily evaluated type hint
   import tensorflow as tf
   ...

How can I achieve the same when defining type aliases?

# types.py  -- a collection of our own type aliases
from typing import TYPE_CHECKING, TypeAlias

if TYPE_CHECKING:
    import tensorflow as tf

MyTensorAlias = "tf.Tensor"  # <-- now this is just a string assignment and will not evaluate to a `ForwardRef`


# Notes of things that do not work:
# 1) annotating with TypeAlias
MyOtherTensorAlias: TypeAlias = "tf.Tensor"  # <-- still just a str and will cause runtime errors

# 2) Aliases behind `TYPE_CHECKING`
if TYPE_CHECKING:
    MyLazyTensorAlias = tf.Tensor  # <-- this is fine for static type checkers, but raises at runtime when `types has no member MyLazyTensorAlias` anymore.

It is possible to achieve lazy evaluation with typing.NewType(), however a subtype is not what I want (because then I need to cast all users, e.g. Subtype(tf.Tensor(...)).

Edit:

The way it fails at runtime is specific to the new | operator. Consider:

Python 3.10.8 (main, Nov 24 2022, 08:09:04) [Clang 14.0.6 ]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.7.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: from typing import TypeAlias
   ...: Alias: TypeAlias = "Original"
   ...: def func(a: Alias | None):
   ...:     pass
   ...:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[1], line 3
      1 from typing import TypeAlias
      2 Alias: TypeAlias = "Original"
----> 3 def func(a: Alias | None):
      4     pass

TypeError: unsupported operand type(s) for |: 'str' and 'NoneType'

In [2]: from typing import TypeAlias,Optional
   ...:
   ...: Alias: TypeAlias = "Original"
   ...: def func(a: Optional[Alias]):
   ...:     pass
   ...:

In [3]: class Original:
   ...:     pass
   ...:

In [4]: def func(a: Original | None):
   ...:     pass
   ...:
1

There are 1 best solutions below

1
Zacharias030 On

As @Jasminj pointed out, MyOtherTensorAlias: TypeAlias = "tf.Tensor" should work, but note that the docs state that type union via X | Y cannot be used with forward references (https://docs.python.org/3/library/stdtypes.html#union-type)