devork

E pur si muove

Inspecting un-imported modules using ast

Sunday, June 24, 2012

or

Skipping modules in py.test using marks but no importing

As I've said before, py.test is my favourite testing tool. One if it's features is that it allows you to mark tests with arbitrary "marks", built in ones are e.g. skipif, xfail etc. But py.test's extension mechanism allows you to easily add behaviour on any other marks you might want.

Recently I've been starting to write some testing code for a Django project (using the great-and-still-improving pytest-django plugin), but I have one strange requirement, on our CI hosts we only want to run the Django tests for a few platforms (We write monitoring software so support way more platforms then is sane). So the natural thing to was mark those test modules as needing Django so we can skip them.

This is how you mark a test module as such, each test written in this module will have the "django" mark applied:

# test_foo.py
import django
import pytest
import foo

pytestmark = pytest.mark.django
# Or, using multiple marks
pytestmark = [pytest.mark.django, py.test.mark.another_mark]

I've also imported django itself, usually test modules need to do this one way or another, at least indirectly via the module they are testing. Now we could easily write py.test conftest file to skip these tests:

# conftest.py
import pytest

ENABLE_DJANGO = True  # imagine some complicated expression/function

def pytest_setup_item(item):
    if 'django' in item.keywords and not ENABLE_DJANGO:
        pytest.skip('Django tests are disabled')

But we have an additional problem, we don't even build Django for all platforms so can't import these test modules in the first place. This means py.test fails while collecting the test modules. There is another hook available in py.test which allows us to skip modules without importing them: pytest_ignore_collect(). But we still need to be able to read the marker inside this module without importing, time to enter ast.

The ast module can parse python code into an abstract syntax tree, which represents the code in a tree of nodes according to the grammar. This is actually a part of the normal compilation process, but just as an internal step. Enough introduction, it was pleasantly easy to use it to find the marks in our test modules:

def _static_mark(path):
    """Return the pytestmark of a test file whithout executing it"""
    module = ast.parse(path.read(), str(path))
    marks = []
    for node in module.body:
        if not isinstance(node, ast.Assign):
            continue
        if len(node.targets) > 1:
            continue
        if not isinstance(node.targets[0], ast.Name):
            continue
        if node.targets[0].id != 'pytestmark':
            continue
        if isinstance(node.value, ast.Attribute):
            marks.append(node.value.attr)
            continue
        if isinstance(node.value, ast.List):
            for elem in node.value.elts:
                if isinstance(elem, ast.Attribute):
                    marks.append(elem.attr)
    return marks

That's it. Essentially we're looking for an ast.Assign node which has a target of pytestmark that is, something is being assigned to pytestmark. Once we find such a node we make sure to only accept a few right-hand-side expressions, namely either pytest.mark.the_mark or [pytest.mark.mark0, pytest.mark.mark1].

Now this is obviously not bulletproof, but it does keep it nice and straight forward. And I thought is a nice example of how to the ast module can be useful.

Subscribe to: Posts (Atom)