Skip to content

Commit 5d9d6ce

Browse files
amurekiamureki
committed
Ref #61 -- Introduce generators for each integer field
1 parent 71d04c8 commit 5d9d6ce

File tree

3 files changed

+131
-19
lines changed

3 files changed

+131
-19
lines changed

‎model_bakery/generators.py‎

Lines changed: 9 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from decimal import Decimal
22
from typing import Any, Callable, Dict, Optional, Type, Union
33

4-
from django.db.backends.base.operations import BaseDatabaseOperations
54
from django.db.models import (
65
AutoField,
76
BigAutoField,
@@ -81,29 +80,20 @@
8180
IntegerRangeField = None
8281

8382

84-
def _make_integer_gen_by_range(field_type: Any) -> Callable:
85-
min_int, max_int = BaseDatabaseOperations.integer_field_ranges[field_type.__name__]
86-
87-
def gen_integer():
88-
return random_gen.gen_integer(min_int=min_int, max_int=max_int)
89-
90-
return gen_integer
91-
92-
9383
default_mapping = {
9484
ForeignKey: random_gen.gen_related,
9585
OneToOneField: random_gen.gen_related,
9686
ManyToManyField: random_gen.gen_m2m,
9787
BooleanField: random_gen.gen_boolean,
98-
AutoField: _make_integer_gen_by_range(AutoField),
99-
BigAutoField: _make_integer_gen_by_range(BigAutoField),
100-
IntegerField: _make_integer_gen_by_range(IntegerField),
101-
SmallAutoField: _make_integer_gen_by_range(SmallAutoField),
102-
BigIntegerField: _make_integer_gen_by_range(BigIntegerField),
103-
SmallIntegerField: _make_integer_gen_by_range(SmallIntegerField),
104-
PositiveBigIntegerField: _make_integer_gen_by_range(PositiveBigIntegerField),
105-
PositiveIntegerField: _make_integer_gen_by_range(PositiveIntegerField),
106-
PositiveSmallIntegerField: _make_integer_gen_by_range(PositiveSmallIntegerField),
88+
AutoField: random_gen.gen_auto_field,
89+
BigAutoField: random_gen.gen_positive_big_integer,
90+
IntegerField: random_gen.gen_regular_integer,
91+
SmallAutoField: random_gen.gen_positive_small_integer,
92+
BigIntegerField: random_gen.gen_big_integer,
93+
SmallIntegerField: random_gen.gen_small_integer,
94+
PositiveBigIntegerField: random_gen.gen_positive_big_integer,
95+
PositiveIntegerField: random_gen.gen_positive_integer,
96+
PositiveSmallIntegerField: random_gen.gen_positive_small_integer,
10797
FloatField: random_gen.gen_float,
10898
DecimalField: random_gen.gen_decimal,
10999
BinaryField: random_gen.gen_byte_string,

‎model_bakery/random_gen.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,9 +77,87 @@ def gen_from_choices(choices: List) -> Callable:
7777

7878

7979
def gen_integer(min_int: int = -MAX_INT, max_int: int = MAX_INT) -> int:
80+
warnings.warn(
81+
"gen_integer() may cause overflow errors with Django integer fields due to "
82+
"large default MAX_INT value. Consider using field-specific generators instead:\n"
83+
"- gen_positive_small_integer() for PositiveSmallIntegerField\n"
84+
"- gen_small_integer() for SmallIntegerField\n"
85+
"- gen_regular_integer() for IntegerField\n"
86+
"- gen_positive_integer() for PositiveIntegerField\n"
87+
"- gen_big_integer() for BigIntegerField\n"
88+
"- gen_positive_big_integer() for PositiveBigIntegerField\n"
89+
"See model_bakery.random_gen documentation for more details.",
90+
DeprecationWarning,
91+
stacklevel=2,
92+
)
8093
return baker_random.randint(min_int, max_int)
8194

8295

96+
def _get_field_range(field_name: str):
97+
"""Get field range from Django's BaseDatabaseOperations."""
98+
from django.db.backends.base.operations import BaseDatabaseOperations
99+
100+
return BaseDatabaseOperations.integer_field_ranges.get(
101+
field_name, (-MAX_INT, MAX_INT)
102+
)
103+
104+
105+
def gen_small_integer(min_int: int = None, max_int: int = None) -> int:
106+
"""Generate integer for SmallIntegerField."""
107+
field_min, field_max = _get_field_range("SmallIntegerField")
108+
actual_min = min_int if min_int is not None else field_min
109+
actual_max = max_int if max_int is not None else field_max
110+
return baker_random.randint(actual_min, actual_max)
111+
112+
113+
def gen_positive_small_integer(min_int: int = None, max_int: int = None) -> int:
114+
"""Generate integer for PositiveSmallIntegerField."""
115+
field_min, field_max = _get_field_range("PositiveSmallIntegerField")
116+
actual_min = min_int if min_int is not None else max(field_min, 1)
117+
actual_max = max_int if max_int is not None else field_max
118+
return baker_random.randint(actual_min, actual_max)
119+
120+
121+
def gen_positive_integer(min_int: int = None, max_int: int = None) -> int:
122+
"""Generate integer for PositiveIntegerField."""
123+
field_min, field_max = _get_field_range("PositiveIntegerField")
124+
actual_min = min_int if min_int is not None else max(field_min, 1)
125+
actual_max = max_int if max_int is not None else field_max
126+
return baker_random.randint(actual_min, actual_max)
127+
128+
129+
def gen_big_integer(min_int: int = None, max_int: int = None) -> int:
130+
"""Generate integer for BigIntegerField."""
131+
field_min, field_max = _get_field_range("BigIntegerField")
132+
actual_min = min_int if min_int is not None else field_min
133+
actual_max = max_int if max_int is not None else field_max
134+
return baker_random.randint(actual_min, actual_max)
135+
136+
137+
def gen_positive_big_integer(min_int: int = None, max_int: int = None) -> int:
138+
"""Generate integer for PositiveBigIntegerField."""
139+
field_min, field_max = _get_field_range("PositiveBigIntegerField")
140+
actual_min = min_int if min_int is not None else max(field_min, 1)
141+
actual_max = max_int if max_int is not None else field_max
142+
return baker_random.randint(actual_min, actual_max)
143+
144+
145+
def gen_regular_integer(min_int: int = None, max_int: int = None) -> int:
146+
"""Generate integer for IntegerField."""
147+
field_min, field_max = _get_field_range("IntegerField")
148+
actual_min = min_int if min_int is not None else field_min
149+
actual_max = max_int if max_int is not None else field_max
150+
return baker_random.randint(actual_min, actual_max)
151+
152+
153+
def gen_auto_field(min_int: int = None, max_int: int = None) -> int:
154+
"""Generate integer for AutoField."""
155+
field_min, field_max = _get_field_range("AutoField")
156+
actual_min = min_int if min_int is not None else max(field_min, 1)
157+
actual_max = max_int if max_int is not None else field_max
158+
return baker_random.randint(actual_min, actual_max)
159+
160+
83161
def gen_float() -> float:
84162
return baker_random.random() * gen_integer()
85163

‎tests/test_baker.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1184,3 +1184,47 @@ def test_make_with_auto_now_and_fill_optional(self):
11841184
assert instance.created == created
11851185
assert instance.updated == updated
11861186
assert instance.sent_date == sent_date
1187+
1188+
1189+
class TestFieldSpecificIntegerGenerators:
1190+
@pytest.mark.django_db
1191+
def test_gen_positive_small_integer_works_safely(self):
1192+
obj = baker.make(
1193+
models.DummyPositiveIntModel,
1194+
positive_small_int_field=random_gen.gen_positive_small_integer(min_int=1),
1195+
)
1196+
1197+
assert 1 <= obj.positive_small_int_field <= 32767
1198+
1199+
def test_field_specific_generators_respect_constraints(self):
1200+
obj = baker.make(
1201+
models.DummyPositiveIntModel,
1202+
positive_small_int_field=random_gen.gen_positive_small_integer(
1203+
min_int=100, max_int=200
1204+
),
1205+
)
1206+
1207+
assert 100 <= obj.positive_small_int_field <= 200
1208+
1209+
def test_gen_integer_shows_deprecation_warning(self):
1210+
import warnings
1211+
1212+
with warnings.catch_warnings(record=True) as w:
1213+
warnings.simplefilter("always")
1214+
1215+
# This should trigger the deprecation warning
1216+
value = random_gen.gen_integer(min_int=1, max_int=100)
1217+
1218+
assert len(w) == 1
1219+
assert issubclass(w[0].category, DeprecationWarning)
1220+
assert "gen_integer() may cause overflow errors" in str(w[0].message)
1221+
assert "gen_positive_small_integer()" in str(w[0].message)
1222+
1223+
assert 1 <= value <= 100
1224+
1225+
def test_automatic_generation_still_works(self):
1226+
obj = baker.make(models.DummyPositiveIntModel)
1227+
1228+
assert 1 <= obj.positive_small_int_field <= 32767
1229+
assert isinstance(obj.positive_int_field, int)
1230+
assert isinstance(obj.positive_big_int_field, int)

0 commit comments

Comments
 (0)