CODEWITHSUNDEEP
pythonmypypython-typing
dspencer
dspencer

Reputation: 4461

How to type hint with an optional import?

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?

Upvotes: 21

Views: 4210

Answers (4)

Ivan Stepanov
Ivan Stepanov

Reputation: 1

Had similar issue, solved by type hinting the package with Any and set it to None before actual import.

from typing import Any

OpenAIResource: Any = None
try:
    from promptchain_ext.openai.resource import OpenAIResource
except ImportError:
    pass

Upvotes: 0

Pablo R. Mier
Pablo R. Mier

Reputation: 737

Another approach that can avoid problems with some linters (e.g Pylance) is this one:

from typing import Any, TYPE_CHECKING

DataFrame = Any

if TYPE_CHECKING:
    try:
        from pandas import DataFrame

    except ImportError:
        pass

DataFrameType = TypeVar("DataFrameType", bound=DataFrame)

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

Upvotes: -1

JustinFisher
JustinFisher

Reputation: 693

Here's the solution I've tentatively been using, which seems to work in PyCharm's type-checker, though I haven't tried MyPy.

from typing import TypeVar, TYPE_CHECKING

PANDAS_CONFIRMED = False
if TYPE_CHECKING:
    try:
        import pandas as pd
        PANDAS_CONFIRMED = True
    except ImportError:
        pass 

if PANDAS_CONFIRMED:
    DataFrameType = pd.DataFrame
else:
    DataFrameType = TypeVar('DataFrameType')

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

This has the advantage of always defining the function (so if someone runs code that calls my_func, they'll get an informative ImportError rather than a misleading AttributeError). This also always offers some sort of type-hint even when pandas is not installed, without trying to import pandas prematurely at runtime. The if-else structure makes PyCharm view some instances of DataFrameType as being Union[DataFrame, DataFrameType] but it still provides linting information that is well-suited for a DataFrame, and in some cases, like my_func's output, it somehow infers that a DataFrameType instance will always be a DataFrame.

Upvotes: 0

Michael0x2a
Michael0x2a

Reputation: 63978

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.

Upvotes: 19

Related Questions