-
Notifications
You must be signed in to change notification settings - Fork 41
/
codegen.py
139 lines (103 loc) · 4.04 KB
/
codegen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
from __future__ import annotations
import atexit
import importlib.util
import itertools
import logging
import pathlib
import shutil
import tempfile
import typing
from dataclasses import field, dataclass
from functools import lru_cache
from types import ModuleType
from typing import Any
import mako.template
from appdirs import AppDirs
from .compat import path_with_stem
cache_dir = pathlib.Path(AppDirs("pytest-factoryboy").user_cache_dir)
logger = logging.getLogger(__name__)
@dataclass
class FixtureDef:
name: str
function_name: typing.Literal["model_fixture", "attr_fixture", "factory_fixture", "subfactory_fixture"]
function_kwargs: dict[str, Any] = field(default_factory=dict)
deps: list[str] = field(default_factory=list)
related: list[str] = field(default_factory=list)
@property
def kwargs_var_name(self):
return f"_{self.name}__kwargs"
module_template = mako.template.Template(
"""\
import pytest
from pytest_factoryboy.fixture import (
attr_fixture,
factory_fixture,
model_fixture,
subfactory_fixture,
)
def _fixture(related):
def fixture_maker(fn):
fn._factoryboy_related = related
return pytest.fixture(fn)
return fixture_maker
% for fixture_def in fixture_defs:
${ fixture_def.kwargs_var_name } = {}
@_fixture(related=${ repr(fixture_def.related) })
def ${ fixture_def.name }(
% for dep in ["request"] + fixture_def.deps:
${ dep },
% endfor
):
return ${ fixture_def.function_name }(request, **${ fixture_def.kwargs_var_name })
% endfor
"""
)
init_py_content = '''\
"""Pytest-factoryboy generated fixtures.
This module and the other modules in this package are automatically generated by
pytest-factoryboy. They will be rewritten on the next run.
"""
'''
@lru_cache() # This way we reuse the same folder for the whole execution of the program
def make_temp_folder(package_name: str) -> pathlib.Path:
"""Create a temporary folder and automatically delete it when the process exit."""
path = pathlib.Path(tempfile.mkdtemp()) / package_name
path.mkdir(parents=True, exist_ok=True)
atexit.register(shutil.rmtree, str(path))
return path
@lru_cache() # This way we reuse the same folder for the whole execution of the program
def create_package(package_name: str, init_py_content=init_py_content) -> pathlib.Path:
path = cache_dir / package_name
try:
if path.exists():
shutil.rmtree(str(path))
path.mkdir(parents=True, exist_ok=False)
except OSError: # Catch cases where the directory can't be removed or can't be created
logger.warning(f"Can't create the cache directory {path}. Using a temporary directory instead.", exc_info=True)
return make_temp_folder(package_name)
(path / "__init__.py").write_text(init_py_content)
return path
def make_module(code: str, module_name: str, package_name: str) -> ModuleType:
tmp_module_path = create_package(package_name) / f"{module_name}.py"
counter = itertools.count(1)
stem = tmp_module_path.stem
while tmp_module_path.exists():
count = next(counter)
new_stem = f"{stem}_{count}"
tmp_module_path = path_with_stem(tmp_module_path, new_stem)
logger.info(f"Writing content of {module_name!r} into {tmp_module_path}.")
tmp_module_path.write_text(code)
name = f"{package_name}.{module_name}"
spec = importlib.util.spec_from_file_location(name, tmp_module_path)
assert spec # NOTE: satisfy `mypy`
mod = importlib.util.module_from_spec(spec)
assert spec.loader # NOTE: satisfy `mypy`
spec.loader.exec_module(mod)
return mod
def make_fixture_model_module(model_name, fixture_defs: list[FixtureDef]):
code = module_template.render(fixture_defs=fixture_defs)
generated_module = make_module(code, module_name=model_name, package_name="_pytest_factoryboy_generated_fixtures")
for fixture_def in fixture_defs:
assert hasattr(generated_module, fixture_def.kwargs_var_name)
setattr(generated_module, fixture_def.kwargs_var_name, fixture_def.function_kwargs)
return generated_module