Annotate a function argument as being a specific module

Question

I have a pytest fixture that imports a specific module. This is needed as importing the module is very expensive, so we don't want to do it on import-time (i.e. during pytest test collection). This results in code like this:

@pytest.fixture
def my_module_fix():
    import my_module
    yield my_module

def test_something(my_module_fix):
    assert my_module_fix.my_func() = 5

I am using PyCharm and would like to have type-checking and autocompletion in my tests. To achieve that, I would somehow have to annotate the my_module_fix parameter as having the type of the my_module module.

I have no idea how to achieve that. All I found is that I can annotate my_module_fix as being of type types.ModuleType, but that is not enough: It is not any module, it is always my_module.

Answer 1

If I get your question, you have two (or three) separate goals

1. Deferred import of slowmodule
2. Autocomplete to continue to work as if it was a standard import
3. (Potentially?) typing (e.g. mypy?) to continue to work

I can think of at least five different approaches, though I'll only briefly mention the last because it's insane.

---

Import the module inside your tests

This is (by far) the most common and IMHO preferred solution.

e.g. instead of

import slowmodule

def test_foo():
    slowmodule.foo()

def test_bar():
    slowmodule.bar()

you'd write:

def test_foo():
    import slowmodule
    slowmodule.foo()

def test_bar():
    import slowmodule
    slowmodule.bar()

• [deferred importing] Here, the module will be imported on-demand/lazily. So if you have pytest setup to fail-fast, and another test fails before pytest gets to your (test_foo, test_bar) tests, the module will never be imported and you'll never incur the runtime cost.

  Because of Python's module cache, subsequent import statements won't actually re-import the module, just grab a reference to the already-imported module.

• [autocomplete/typing] Of course, autocomplete will continue to work as you expect in this case. This is a perfectly fine import pattern.

While it does require adding potentially many additional import statements (one inside each test function), it's immediately clear what is going on (regardless of whether it's clear why it's going on).

---

[3.7+] Proxy your module with module __getattr__

If you create a module (e.g. slowmodule_proxy.py) with the contents like:

def __getattr__(name):
    import slowmodule
    return getattr(slowmodule, name)

And in your tests, e.g.

import slowmodule

def test_foo():
    slowmodule.foo()

def test_bar():
    slowmodule.bar()

instead of:

import slowmodule

you write:

import slowmodule_proxy as slowmodule

• [deferred import] Thanks to PEP-562, you can "request" any name from slowmodule_proxy and it will fetch and return the corresponding name from slowmodule. Just as above, including the import inside the function will cause slowmodule to be imported only when the function is called and executed instead of on module load. Module caching still applies here of course, so you're only incurring the import penalty once per interpreter session.

• [autocomplete] However, while deferred importing will work (and your tests run without issue), this approach (as stated so far) will "break" autocomplete:

Now we're in the realm of PyCharm. Some IDEs will perform "live" analysis of modules and actually load up the module and inspect its members. (PyDev had this option). If PyCharm did this, implementing module.__dir__ (same PEP) or __all__ would allow your proxy module to masquerade as the actual slowmodule and autocomplete would work.† But, PyCharm does not do this.

Nonetheless, you can fool PyCharm into giving you autocomplete suggestions:

if False:
    import slowmodule
else:
    import slowmodule_proxy as slowmodule

The interpreter will only execute the else branch, importing the proxy and naming it slowmodule (so your test code can continue to reference slowmodule unchanged).

But PyCharm will now provide autocompletion for the underlying module:

† _{While live-analysis can be an incredibly helpful, there's also a (potential) security concern that comes with it that static syntax analysis doesn't have. And the maturation of type hinting and stub files has made it less of an issue still.}

---

Proxy slowmodule explicitly

If you really hated the dynamic proxy approach (or the fact that you have to fool PyCharm in this way), you could proxy the module explicitly.

(You'd likely only want to consider this if the slowmodule API is stable.)

If slowmodule has methods foo and bar you'd create a proxy module like:

def foo(*args, **kwargs):
    import slowmodule
    return slowmodule.foo(*args, **kwargs)

def bar(*args, **kwargs):
    import slowmodule
    return slowmodule.bar(*args, **kwargs)

(Using args and kwargs to pass arguments through to the underlying callables. And you could add type hinting to these functions to mirror the slowmodule functions.)

And in your test,

import slowmodule_proxy as slowmodule

Same as before. Importing inside the method gives you the deferred importing you want and the module cache takes care of multiple import calls.

And since it's a real module whose contents can be statically analyzed, there's no need to "fool" PyCharm.

So the benefit of this solution is that you don't have a bizarre looking if False in your test imports. This, however, comes at the (substantial) cost of having to maintain a proxy file alongside your module -- which could prove painful in the case that slowmodule's API wasn't stable.

---

[3.5+] Use importlib's LazyLoader instead of a proxy module

Instead of the proxy module slowmodule_proxy, you could follow a pattern similar to the one shown in the importlib docs

>>> import importlib.util
>>> import sys
>>> def lazy_import(name):
...     spec = importlib.util.find_spec(name)
...     loader = importlib.util.LazyLoader(spec.loader)
...     spec.loader = loader
...     module = importlib.util.module_from_spec(spec)
...     sys.modules[name] = module
...     loader.exec_module(module)
...     return module
...
>>> lazy_typing = lazy_import("typing")
>>> #lazy_typing is a real module object,
>>> #but it is not loaded in memory yet.

You'd still need to fool PyCharm though, so something like:

if False:
    import slowmodule
else:
    slowmodule = lazy_import('slowmodule')

would be necessary.

Outside of the single additional level of indirection on module member access (and the two minor version availability difference), it's not immediately clear to me what, if anything, there is to be gained from this approach over the previous proxy module method, however.

---

Use importlib's Finder/Loader machinery to hook import (don't do this)

You could create a custom module Finder/Loader that would (only) hook your slowmodule import and, instead load, for example your proxy module.

Then you could just import that "importhook" module before you imported slowmode in your tests, e.g.

import myimporthooks
import slowmodule

def test_foo():
    ...

(Here, myimporthooks would use importlib's finder and loader machinery to do something simlar to the importhook package but intercept and redirect the import attempt rather than just serving as an import callback.)

But this is crazy. Not only is what you want (seemingly) achievable through (infinitely) more common and supported methods, but it's incredibly fragile, error-prone and, without diving into the internals of PyTest (which may mess with module loaders itself), it's hard to say whether it'd even work.

Answer 2

When Pytest collects files to be tested, modules are only imported once, even if the same import statement appears in multiple files.

To observe when my_module is imported, add a print statement and then use the Pytest -s flag (short for --capture=no), to ensure that all standard output is displayed.

my_module.py

answer: int = 42
print("MODULE IMPORTED: my_module.py")

You could then add your test fixture to a conftest.py file:

conftest.py

import pytest


@pytest.fixture
def my_module_fix():
    import my_module
    yield my_module

Then in your test files, my_module.py may be imported to add type hints:

test_file_01.py

import my_module


def test_something(my_module_fix: my_module):
    assert my_module_fix.answer == 42

test_file_02.py

import my_module


def test_something2(my_module_fix: my_module):
    assert my_module_fix.answer == 42

Then run Pytest to display all standard output and verify that the module is only imported once at runtime.

pytest -s ./

Output from Pytest

platform linux -- Python 3.9.7, pytest-6.2.5
rootdir: /home/your_username/your_repo
collecting ... MODULE IMPORTED: my_module.py  <--- Print statement executed once
collected 2 items

test_file_01.py .
test_file_02.py .

Answer 3

This is quick and naive, but can you possibly annotate your tests with the "quote-style" annotation that is meant for different purposes, but may suit here by skipping the import at runtime but still help your editor?

def test_something(my_module_fix: "my_module"):

In a quick test, this seems to accomplish it at least for my setup.

Answer 4

Although it might not be considered a 'best practice', to keep your specific use case simple, you could just lazily import the module directly in your test where you need it.

def test_something():
    import my_module
    assert my_module.my_func() = 5

I believe Pytest will only import the module when the applicable tests run. Pytest should also 'cache' the import as well so that if multiple tests import the module, it is only actually imported once. This may also solve your autocomplete issues for your editor.

Side-note: Avoid writing code in a specific way to cater for a specific editor. Keep it simple, not everyone who looks at your code will use Pycharm.

Answer 5

would like to have type-checking and autocompletion in my tests

It sounds like you want Something that fits as the type symbol of your test function:

def test_something(my_module_fix: Something):
    assert my_module_fix.my_func() = 5

... and from this, [hopefully] your IDE can make some inferences about my_module_fix. I'm not a PyCharm user, so I can't speak to what it can tell you from type signatures, but I can say this isn't something that's readily available.

For some intuition, in this example -- Something is a ModuleType like 3 is an int. Analogously, accessing a nonexistent attribute of Something is like doing something not allowed with 3. Perhaps, like accessing an attribute of it (3.__name__).

But really, I seems like you're thinking about this from the wrong direction. The question a type signature answers is: what contract must this [these] argument[s] satisfy for this function to use it [them]. Using the example above, a type of 3 is the kind of thing that is too specific to make useful functions:

def add_to(i: 3):
    return i

Perhaps a better name for your type is:

def test_something(my_module_fix: SomethingThatHasMyFuncMethod):
    assert my_module_fix.my_func() = 5

So the type you probably want is something like this:

class SomethingThatHasMyFuncMethod:
   def my_func(self) -> int: ... 

You'll need to define it (maybe in a .pyi file). See here for info.

---

Finally, here's some unsolicited advice regarding:

importing the module is very expensive, so we don't want to do it on import-time

You should probably employ some method of making this module do it's thing lazily. Some of the utils in the django framework could serve as a reference point. There's also the descriptor protocol which is a bit harder to grok but may suit your needs.

Answer 6

You want to add type hinting for the arguments of test_something, in particular to my_module_fix. This question is hard because we can't import my_module at the top.

Well, what is the type of my_module? I'm going to assume that if you do import my_module and then type(my_module) you will get <class 'module'>.

If you're okay with not telling users exactly which module it has to be, then you could try this.

from types import ModuleType

@pytest.fixture
def my_module_fix():
    import my_module
    yield my_module

def test_something(my_module_fix: ModuleType):
    assert my_module_fix.my_func() = 5

The downside is that a user might infer that any old module would be suitable for my_module_fix.

You want it to be more specific? There's a cost to that, but we can get around some of it (it's hard to speak on performance in VS Code vs PyCharm vs others).

In that case, let's use TYPE_CHECKING. I think this was introduced in Python 3.6.

from types import ModuleType
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import my_module

@pytest.fixture
def my_module_fix():
    import my_module
    yield my_module

def test_something(my_module_fix: my_module):
    assert my_module_fix.my_func() = 5

You won't pay these costs during normal run time. You will pay the cost of importing my_module when your type checker is doing what it does.

If you're on an earlier verison of Python, you might need to do this instead.

from types import ModuleType
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import my_module

@pytest.fixture
def my_module_fix():
    import my_module
    yield my_module

def test_something(my_module_fix: "my_module"):
    assert my_module_fix.my_func() = 5

The type checker in VS Code is smart enough to know what's going on here. I can't speak for other type checkers.