How to type hint with an optional import?

Question

When using an optional import, i.e. the package is only imported inside a function as I want it to be an optional dependency of my package, is there a way to type hint the return type of the function as one of the classes belonging to this optional dependency?

To give a simple example with pandas as an optional dependency:

def my_func() -> pd.DataFrame:                                                  
    import pandas as pd                                                         
    return pd.DataFrame()                                                       

df = my_func()

In this case, since the import statement is within my_func, this code will, not surprisingly, raise:

NameError: name 'pd' is not defined

If the string literal type hint were used instead, i.e.:

def my_func() -> 'pd.DataFrame':                                                
    import pandas as pd                                                         
    return pd.DataFrame()                                                       

df = my_func()

the module can now be executed without issue, but mypy will complain:

error: Name 'pd' is not defined

How can I make the module execute successfully and retain the static type checking capability, while also having this import be optional?

Answer 1

Try sticking your import inside of an if typing.TYPE_CHECKING statement at the top of your file. This variable is always false at runtime but is treated as always true for the purposes of type hinting.

For example:

# Lets us avoid needing to use forward references everywhere
# for Python 3.7+
from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import pandas as pd

def my_func() -> pd.DataFrame:  
    import pandas as pd                                                 
    return pd.DataFrame()

You can also do if False:, but I think that makes it a little harder for somebody to tell what's going on.

One caveat is that this does mean that while pandas will be an optional dependency at runtime, it'll still be a mandatory one for the purposes of type checking.

Another option you can explore using is mypy's --always-true and --always-false flags. This would give you finer-grained control over which parts of your code are typechecked. For example, you could do something like this:

try:
    import pandas as pd
    PANDAS_EXISTS = True
except ImportError:
    PANDAS_EXISTS = False

if PANDAS_EXISTS:
    def my_func() -> pd.DataFrame:                                                   
        return pd.DataFrame()

...then do mypy --always-true=PANDAS_EXISTS your_code.py to type check it assuming pandas is imported and mypy --always-false=PANDAS_EXISTS your_code.py to type check assuming it's missing.

This could help you catch cases where you accidentally use a function that requires pandas from a function that isn't supposed to need it -- though the caveats are that (a) this is a mypy-only solution and (b) having functions that only sometimes exist in your library might be confusing for the end-user.