How to register_assert_rewrite on relative imports

[AckslD]

Consider the two following scenarios:

Non-relative import

Lets say I have the following files

test.py
utils.py

# test.py
import pytest
pytest.register_assert_rewrite('utils')
from utils import assert_equal


def test_equal():
    a = [0] * 100
    b = [1] * 100
    assert_equal(a, b)

# utils
def assert_equal(a, b):
    assert a == b

Then running pytest test.py gives nice assert output:

a = [0, 0, 0, 0, 0, 0, ...], b = [1, 1, 1, 1, 1, 1, ...]

    def assert_equal(a, b):
>       assert a == b
E       assert [0, 0, 0, 0, 0, 0, ...] == [1, 1, 1, 1, 1, 1, ...]
E         At index 0 diff: 0 != 1
E         Use -v to get the full diff

but in the next scenario:

Relative import

Lets say I have the following files

__init__.py
test.py
utils.py

# test.py
import pytest
pytest.register_assert_rewrite('utils')
from .utils import assert_equal


def test_equal():
    a = [0] * 100
    b = [1] * 100
    assert_equal(a, b)

# utils
def assert_equal(a, b):
    assert a == b

Differnce to the first scenario is:

Only in second scenario: __init__.py
4c3
< from utils import assert_equal
---
> from .utils import assert_equal

Then running pytest test.py only gives.

a = [0, 0, 0, 0, 0, 0, ...], b = [1, 1, 1, 1, 1, 1, ...]

    def assert_equal(a, b):
>       assert a == b
E       AssertionError

ie the rewrite did not work.

Any idea how to do this correctly?

[nicoddemus]

I'm not sure register_assert_rewrite will work for relative imports, I think pytest will only "see" the absolute import when hooked to the Python import machinery.

You can use --trace to see the internal logging, which should show which imports pytest is attempting to match with the ones given to register_assert_rewrite.

[nicoddemus]

Argh sorry, I meant --trace-config --debug.

EDIT: I just tested it and it doesn't really print the information I was hoping for, it only prints when there's a match.

Here is the code responsible to decide if a module should be rewritten or not, in case you want to dig into it:

pytest/src/_pytest/assertion/rewrite.py

Lines 232 to 243 in 6d6bc97

	def _is_marked_for_rewrite(self, name: str, state: "AssertionState") -> bool:
	try:
	return self._marked_for_rewrite_cache[name]
	except KeyError:
	for marked in self._must_rewrite:
	if name == marked or name.startswith(marked + "."):
	state.trace(f"matched marked file {name!r} (from {marked!r})")
	self._marked_for_rewrite_cache[name] = True
	return True
	self._marked_for_rewrite_cache[name] = False
	return False

[AckslD]

I tried to see what's going on here. I can see that in the case of not relative imports the function find_spec is called but not in the case with the relative import. But I'm not sure why that is or why it's called in one case and the other not.

[nicoddemus]

@asottile do you know if find_spec works differently for relative imports?

[asottile]

it does and it doesn't -- so what's happening here is because you're running pytest t.py and you've got an __init__.py the name utils is available from two module names, just blank "utils" and the full module path (whatever folder (you haven't shown it)).utils

with an __init__.py there, I'd argue that the blank "utils" module name is incorrect and you should use the fully-qualified name

in fact if I set up a directory structure like:

wat/
    __init__.py
    tests.py
    utils.py

similar to yours except I use register_assert_rewrite('wat.utils') then everything works as I expect

there's an earlier part in import machinery that translates relative imports plus __path__ to absolute imports

[AckslD]

Thanks for the explanation @asottile and sorry for the late reply! This is very useful 👍